from __future__ import annotations

import json
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any


@dataclass
class StimulusSynthesisConfig:
    encoder_model_id: str = "epfl-neuroai/vjepa2-encoder-basic"
    encoder_call: str = "predict_fmri"
    text_to_image_model_id: str = "stabilityai/sdxl-turbo"
    image_to_video_model_id: str = "Lightricks/LTX-Video-0.9.8-13B-distilled"
    default_objective: str = "indices_mean"
    default_device: str = "cuda"
    default_population_size: int = 20
    default_image_max_evals: int = 400
    default_video_max_evals: int = 200
    default_score_frames: int = 24
    default_score_size: int = 224
    default_image_batch_size: int = 16
    default_video_batch_size: int = 8
    default_mutation_rate: float = 0.25
    default_elite_frac: float = 0.35
    default_score_transform: dict[str, Any] | None = field(
        default_factory=lambda: {
            "enabled": False,
            "crop_scale": 0.80,
            "gaussian_sigma": 0.10,
            "num_draws": 4,
            "aggregate": "mean",
            "seed": 0,
        }
    )
    default_image_kwargs: dict[str, Any] = field(
        default_factory=lambda: {"num_inference_steps": 1, "guidance_scale": 0.0, "height": 512, "width": 512}
    )
    default_video_kwargs: dict[str, Any] = field(
        default_factory=lambda: {"num_inference_steps": 8, "num_frames": 25, "height": 512, "width": 512}
    )

    @classmethod
    def from_json_file(cls, path: str | Path) -> "StimulusSynthesisConfig":
        with open(path, "r") as f:
            data = json.load(f)
        return cls(**data)

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "StimulusSynthesisConfig":
        known = {field.name for field in cls.__dataclass_fields__.values()}
        return cls(**{k: v for k, v in data.items() if k in known})

    def to_dict(self) -> dict[str, Any]:
        return {
            "encoder_model_id": self.encoder_model_id,
            "encoder_call": self.encoder_call,
            "text_to_image_model_id": self.text_to_image_model_id,
            "image_to_video_model_id": self.image_to_video_model_id,
            "default_objective": self.default_objective,
            "default_device": self.default_device,
            "default_population_size": self.default_population_size,
            "default_image_max_evals": self.default_image_max_evals,
            "default_video_max_evals": self.default_video_max_evals,
            "default_score_frames": self.default_score_frames,
            "default_score_size": self.default_score_size,
            "default_image_batch_size": self.default_image_batch_size,
            "default_video_batch_size": self.default_video_batch_size,
            "default_mutation_rate": self.default_mutation_rate,
            "default_elite_frac": self.default_elite_frac,
            "default_score_transform": self.default_score_transform,
            "default_image_kwargs": self.default_image_kwargs,
            "default_video_kwargs": self.default_video_kwargs,
        }