# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# SPDX-License-Identifier: Apache-2.0

import json
from dataclasses import asdict, dataclass, field
from typing import Any, Dict, List, Optional, Tuple, Union


@dataclass
class BaseConfig:
    def get(self, attribute_name, default=None):
        return getattr(self, attribute_name, default)

    def pop(self, attribute_name, default=None):
        if hasattr(self, attribute_name):
            value = getattr(self, attribute_name)
            delattr(self, attribute_name)
            return value
        else:
            return default

    def __str__(self):
        return json.dumps(asdict(self), indent=4)


@dataclass
class DataConfig(BaseConfig):
    data_dir: List[str] = field(default_factory=list)
    caption_proportion: Dict[str, int] = field(default_factory=lambda: {"prompt": 1})
    external_caption_suffixes: List[str] = field(default_factory=list)
    external_clipscore_suffixes: List[str] = field(default_factory=list)
    caption_selection_type: str = (
        "clipscore"  # clipscore: use $external_clipscore_suffixes, proportion: use $caption_proportion
    )
    clip_thr_temperature: float = 1.0
    clip_thr: float = 0.0
    del_img_clip_thr: float = 0.0
    sort_dataset: bool = False
    load_text_feat: bool = False
    load_vae_feat: bool = False
    aspect_ratio_type: str = "ASPECT_RATIO_1024"
    transform: str = "default_train"
    type: str = "SanaWebDatasetMS"
    image_size: int = 512
    hq_only: bool = False
    valid_num: int = 0
    data: Any = None
    num_frames: int = 81
    extra: Any = None


@dataclass
class VideoDataConfig(DataConfig):
    data_dir: Dict[str, str] = field(default_factory=lambda: {"video_toy_data: data/video_toy_data"})
    aspect_ratio_type: str = "ASPECT_RATIO_VIDEO_256_MS"
    external_data_filter: Dict[str, Dict[str, Dict[str, float]]] = field(default_factory=lambda: {})
    motion_score_file_thres: Dict[str, Optional[float]] = field(default_factory=dict)
    motion_score_cal_type: str = "average"  # average, max
    target_fps: int = 16
    resample_fps: bool = True
    shuffle_dataset: bool = False
    vae_cache_dir: Optional[str] = None
    json_cache_dir: Optional[str] = None
    load_first_frame: bool = False


@dataclass
class ModelConfig(BaseConfig):
    model: str = "SanaMS_600M_P1_D28"
    teacher: Optional[str] = None
    image_size: int = 512
    mixed_precision: str = "fp16"  # ['fp16', 'fp32', 'bf16']
    fp32_attention: bool = True
    load_from: Optional[str] = None
    discriminator_model: Optional[str] = None
    teacher_model: Optional[str] = None
    teacher_model_weight_dtype: Optional[str] = None
    resume_from: Optional[Union[Dict[str, Any], str]] = field(
        default_factory=lambda: {
            "checkpoint": None,
            "load_ema": False,
            "resume_lr_scheduler": True,
            "resume_optimizer": True,
        }
    )
    aspect_ratio_type: Optional[str] = None
    multi_scale: bool = True
    pe_interpolation: float = 1.0
    micro_condition: bool = False
    attn_type: str = "linear"
    autocast_linear_attn: bool = False
    ffn_type: str = "glumbconv"
    mlp_acts: List[Optional[str]] = field(default_factory=lambda: ["silu", "silu", None])
    mlp_ratio: float = 2.5
    use_pe: bool = False
    pos_embed_type: str = "sincos"  # "sincos", "flux_rope", "wan_rope"
    qk_norm: bool = False
    class_dropout_prob: float = 0.0
    linear_head_dim: int = 32
    cross_norm: bool = False
    cross_attn_type: str = "flash"
    logvar: bool = False
    cfg_scale: int = 4
    cfg_embed: bool = False
    cfg_embed_scale: float = 1.0
    guidance_type: str = "classifier-free"
    pag_applied_layers: List[int] = field(default_factory=lambda: [8])
    # for ladd
    ladd_multi_scale: bool = True
    head_block_ids: Optional[List[int]] = None
    extra: Any = None


@dataclass
class ModelVideoConfig(ModelConfig):
    # stage1
    remove_state_dict_keys: Optional[List[str]] = None
    # stage2
    rope_fhw_dim: Optional[Tuple[int, int, int]] = None
    t_kernel_size: int = 3
    flash_attn_window_count: Optional[List[int]] = None
    pack_latents: bool = False
    encode_image_prompt_embeds: bool = False
    # stage3
    cross_attn_image_embeds: bool = False
    image_latent_mode: str = "video_zero"
    # chunkcasual
    chunk_index: Optional[List[int]] = None


@dataclass
class AEConfig(BaseConfig):
    vae_type: str = "AutoencoderDC"
    vae_pretrained: str = "mit-han-lab/dc-ae-f32c32-sana-1.1-diffusers"
    weight_dtype: str = "float32"
    scale_factor: float = 0.41407
    scaling_factor: Optional[Union[float, List[float]]] = None  # for st-dc-ae
    vae_latent_dim: int = 32
    vae_downsample_rate: int = 32
    sample_posterior: bool = True
    vae_stride: Optional[List[int]] = None
    if_cache: bool = False
    cache_dir: Optional[str] = None
    # Framewise / tiling fields used by LTX2VAE_diffusers for long-video decode.
    use_framewise_encoding: bool = False
    use_framewise_decoding: bool = False
    tile_sample_stride_num_frames: int = 64
    tile_sample_min_num_frames: int = 96
    extra: Any = None


@dataclass
class TextEncoderConfig(BaseConfig):
    text_encoder_name: str = "gemma-2-2b-it"
    caption_channels: int = 2304
    y_norm: bool = True
    y_norm_scale_factor: float = 1.0
    model_max_length: int = 300
    chi_prompt: List[Optional[str]] = field(default_factory=lambda: [])
    extra: Any = None


@dataclass
class ImageEncoderConfig(BaseConfig):
    image_encoder_name: Optional[str] = None
    image_encoder_path: Optional[str] = None
    weight_dtype: Optional[str] = "bf16"


@dataclass
class SchedulerConfig(BaseConfig):
    train_sampling_steps: int = 1000
    predict_flow_v: bool = True
    noise_schedule: str = "linear_flow"
    pred_sigma: bool = False
    learn_sigma: bool = True
    vis_sampler: str = "flow_dpm-solver"
    flow_shift: float = 1.0
    inference_flow_shift: Optional[float] = None
    # logit-normal timestep
    weighting_scheme: Optional[str] = "logit_normal"
    weighting_scheme_discriminator: Optional[str] = "logit_normal_trigflow"
    add_noise_timesteps: List[float] = field(default_factory=lambda: [1.57080])
    logit_mean: float = 0.0
    logit_std: float = 1.0
    logit_mean_discriminator: float = 0.0
    logit_std_discriminator: float = 1.0
    mode_scale: float = 1.29
    sigma_data: float = 1.0
    p_low: Optional[float] = None
    p_high: Optional[float] = None
    timestep_norm_scale_factor: float = 1.0
    pretrain_timestep_norm_scale_factor: float = 1.0
    discrete_norm_timestep: bool = False
    extra: Any = None


@dataclass
class TrainingConfig(BaseConfig):
    num_workers: int = 4
    seed: int = 42
    train_batch_size: int = 32
    train_batch_size_image: int = 32
    early_stop_hours: float = 100
    num_epochs: int = 100
    gradient_accumulation_steps: int = 1
    grad_checkpointing: bool = False
    gradient_clip: float = 1.0
    gc_step: int = 1
    optimizer: Dict[str, Any] = field(
        default_factory=lambda: {"eps": 1.0e-10, "lr": 0.0001, "type": "AdamW", "weight_decay": 0.03}
    )
    optimizer_D: Dict[str, Any] = field(
        default_factory=lambda: {"eps": 1.0e-10, "lr": 0.0001, "type": "AdamW", "weight_decay": 0.03}
    )
    load_from_optimizer: bool = False
    load_from_lr_scheduler: bool = False
    resume_lr_scheduler: bool = True
    lr_schedule: str = "constant"
    lr_schedule_args: Dict[str, int] = field(default_factory=lambda: {"num_warmup_steps": 500})
    auto_lr: Optional[Dict[str, str]] = field(default_factory=lambda: {"rule": "sqrt"})
    ema_rate: float = 0.9999
    eval_batch_size: int = 16
    use_fsdp: bool = False
    fsdp_version: int = 1
    cp_size: int = 0
    use_flash_attn: bool = False
    eval_sampling_steps: int = 250
    lora_rank: int = 4
    log_interval: int = 50
    mask_type: str = "null"
    mask_loss_coef: float = 0.0
    load_mask_index: bool = False
    snr_loss: bool = False
    real_prompt_ratio: float = 1.0
    save_image_epochs: int = 1
    save_model_epochs: int = 1
    save_model_steps: int = 1000000
    visualize: bool = False
    null_embed_root: str = "output/pretrained_models/"
    valid_prompt_embed_root: str = "output/tmp_embed/"
    validation_prompts: List[str] = field(
        default_factory=lambda: [
            "dog",
            "portrait photo of a girl, photograph, highly detailed face, depth of field",
            "Self-portrait oil painting, a beautiful cyborg with golden hair, 8k",
            "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
            "A photo of beautiful mountain with realistic sunset and blue lake, highly detailed, masterpiece",
        ]
    )
    local_save_vis: bool = False
    deterministic_validation: bool = True
    online_metric: bool = False
    eval_metric_step: int = 5000
    online_metric_dir: str = "metric_helper"
    work_dir: str = "/cache/exps/"
    skip_step: int = 0
    loss_type: str = "huber"
    huber_c: float = 0.001
    num_ddim_timesteps: int = 50
    w_max: float = 15.0
    w_min: float = 3.0
    ema_decay: float = 0.95
    debug_nan: bool = False
    ema_update: bool = False
    ema_rate: float = 0.9999
    weight_loss: bool = True
    tangent_warmup_steps: int = 10000
    scm_cfg_scale: Union[float, List[float]] = field(default_factory=lambda: [1.0])
    cfg_interval: Optional[List[float]] = None
    scm_logvar_loss: bool = True
    norm_invariant_to_spatial_dim: bool = True
    norm_same_as_512_scale: bool = False
    g_norm_constant: float = 0.1
    g_norm_r: float = 1.0
    show_gradient: bool = False
    lr_scale: Optional[Dict[str, List[str]]] = None
    # for ladd
    adv_lambda: float = 1.0
    scm_loss: bool = True
    scm_lambda: float = 1.0
    loss_scale: float = 1.0
    r1_penalty: bool = False
    r1_penalty_weight: float = 1.0e-5
    diff_timesteps_D: bool = True
    # for adversarial loss
    suffix_checkpoints: Optional[str] = "disc"
    misaligned_pairs_D: bool = False
    discriminator_loss: str = "cross entropy"
    largest_timestep: float = 1.57080
    train_largest_timestep: bool = False
    largest_timestep_prob: float = 0.5
    extra: Any = None


@dataclass
class TrainVideoConfig(TrainingConfig):
    validation_images: Optional[List[str]] = None
    image_prior_type: Optional[str] = None  # [flux-siglip2
    joint_training_interval: int = 50
    timestep_weight: bool = False
    noise_multiplier: Optional[float] = 0.0
    ltx_image_condition_prob: float = 0.0  # for ltx, the image condition is used for the first frame
    chunk_sampling_strategy: str = "uniform"  # uniform, incremental
    same_timestep_prob: float = (
        0.0  # for incremental sampling, the probability of using the same timestep for all chunks
    )
    # temporal coherence loss for video training
    temporal_coherence_loss: bool = False
    temporal_coherence_weight: float = 0.0


@dataclass
class ControlNetConfig(BaseConfig):
    control_signal_type: str = "scribble"
    validation_scribble_maps: List[str] = field(
        default_factory=lambda: [
            "output/tmp_embed/controlnet/dog_scribble_thickness_3.jpg",
            "output/tmp_embed/controlnet/girl_scribble_thickness_3.jpg",
            "output/tmp_embed/controlnet/cyborg_scribble_thickness_3.jpg",
            "output/tmp_embed/controlnet/Astronaut_scribble_thickness_3.jpg",
            "output/tmp_embed/controlnet/mountain_scribble_thickness_3.jpg",
        ]
    )


@dataclass
class ModelGrowthConfig(BaseConfig):
    """Model growth configuration for initializing larger models from smaller ones"""

    pretrained_ckpt_path: str = ""
    init_strategy: str = "constant"  # ['cyclic', 'block_expand', 'progressive', 'interpolation', 'random', 'constant']
    init_params: Dict[str, Any] = field(
        default_factory=lambda: {
            "expand_ratio": 3,
            "noise_scale": 0.01,
        }
    )
    source_num_layers: int = 20
    target_num_layers: int = 30
    extra: Any = None


@dataclass
class SanaConfig(BaseConfig):
    data: DataConfig
    model: ModelConfig
    vae: AEConfig
    text_encoder: TextEncoderConfig
    scheduler: SchedulerConfig
    train: TrainingConfig
    controlnet: Optional[ControlNetConfig] = None
    model_growth: Optional[ModelGrowthConfig] = None
    work_dir: str = "output/"
    resume_from: Optional[str] = None
    load_from: Optional[str] = None
    debug: bool = False
    caching: bool = False
    report_to: str = "wandb"
    tracker_project_name: str = "sana-video-baseline"
    name: str = "baseline"
    loss_report_name: str = "loss"


@dataclass
class WanTextEncoderConfig(BaseConfig):
    t5_model: str = "umt5_xxl"
    t5_dtype: str = "bfloat16"
    text_len: int = 512
    t5_checkpoint: str = "checkpoints/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth"
    t5_tokenizer: str = "google/umt5-xxl"
    extra: Any = None
    caption_channels: int = 4096


@dataclass
class DistillConfig(BaseConfig):
    pass


@dataclass
class SanaVideoConfig(BaseConfig):
    data: VideoDataConfig
    model: ModelVideoConfig
    vae: AEConfig
    text_encoder: TextEncoderConfig
    scheduler: SchedulerConfig
    train: TrainVideoConfig
    image_data: Optional[DataConfig] = None
    image_encoder: Optional[ImageEncoderConfig] = field(default_factory=lambda: {})
    model_growth: Optional[ModelGrowthConfig] = None
    text_encoder_wan: Optional[WanTextEncoderConfig] = None
    work_dir: str = "output/"
    resume_from: Optional[str] = None
    load_from: Optional[str] = None
    debug: bool = False
    caching: bool = False
    report_to: str = "wandb"
    tracker_project_name: str = "sana-video"
    name: str = "baseline"
    loss_report_name: str = "loss"
    task: str = "t2v"  # t2v or ti2v
    distill: Optional[DistillConfig] = None


@dataclass
class SanaVideoStage1Config(BaseConfig):
    data: DataConfig
    model: ModelVideoConfig
    vae: AEConfig
    text_encoder: TextEncoderConfig
    scheduler: SchedulerConfig
    train: TrainVideoConfig
    model_growth: Optional[ModelGrowthConfig] = None
    work_dir: str = "output/"
    resume_from: Optional[str] = None
    load_from: Optional[str] = None
    debug: bool = False
    caching: bool = False
    report_to: str = "wandb"
    tracker_project_name: str = "sana-video"
    name: str = "baseline"
    loss_report_name: str = "loss"
    task: str = "t2v"  # t2v or ti2v or df


def model_init_config(config: SanaConfig, latent_size: int = 32):

    pred_sigma = getattr(config.scheduler, "pred_sigma", True)
    learn_sigma = getattr(config.scheduler, "learn_sigma", True) and pred_sigma
    return {
        "input_size": latent_size,
        "pe_interpolation": config.model.pe_interpolation,
        "config": config,
        "model_max_length": config.text_encoder.model_max_length,
        "qk_norm": config.model.qk_norm,
        "micro_condition": config.model.micro_condition,
        "caption_channels": config.text_encoder.caption_channels,
        "class_dropout_prob": config.model.class_dropout_prob,
        "y_norm": config.text_encoder.y_norm,
        "attn_type": config.model.attn_type,
        "ffn_type": config.model.ffn_type,
        "mlp_ratio": config.model.mlp_ratio,
        "mlp_acts": list(config.model.mlp_acts),
        "in_channels": config.vae.vae_latent_dim,
        "y_norm_scale_factor": config.text_encoder.y_norm_scale_factor,
        "use_pe": config.model.use_pe,
        "pos_embed_type": config.model.pos_embed_type,
        "linear_head_dim": config.model.linear_head_dim,
        "pred_sigma": pred_sigma,
        "learn_sigma": learn_sigma,
        "cross_norm": config.model.cross_norm,
        "cross_attn_type": config.model.cross_attn_type,
        "timestep_norm_scale_factor": config.scheduler.timestep_norm_scale_factor,
        "discrete_norm_timestep": config.scheduler.discrete_norm_timestep,
    }


def model_video_init_config(config: SanaVideoConfig, latent_size: int = 32):

    pred_sigma = getattr(config.scheduler, "pred_sigma", True)
    learn_sigma = getattr(config.scheduler, "learn_sigma", True) and pred_sigma
    return {
        "input_size": latent_size,
        "pe_interpolation": config.model.pe_interpolation,
        "config": config,
        "model_max_length": config.text_encoder.model_max_length,
        "qk_norm": config.model.qk_norm,
        "micro_condition": config.model.micro_condition,
        "caption_channels": config.text_encoder.caption_channels,
        "class_dropout_prob": config.model.class_dropout_prob,
        "y_norm": config.text_encoder.y_norm,
        "attn_type": config.model.attn_type,
        "ffn_type": config.model.ffn_type,
        "mlp_ratio": config.model.mlp_ratio,
        "mlp_acts": list(config.model.mlp_acts),
        "in_channels": config.vae.vae_latent_dim,
        "use_pe": config.model.use_pe,
        "pos_embed_type": config.model.pos_embed_type,
        "rope_fhw_dim": config.model.rope_fhw_dim,
        "linear_head_dim": config.model.linear_head_dim,
        "pred_sigma": pred_sigma,
        "learn_sigma": learn_sigma,
        "cross_norm": config.model.cross_norm,
        "cross_attn_type": config.model.cross_attn_type,
        "cross_attn_image_embeds": config.model.cross_attn_image_embeds,
        "t_kernel_size": config.model.t_kernel_size,
        "flash_attn_window_count": config.model.flash_attn_window_count,
        "pack_latents": config.model.pack_latents,
    }