# 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

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

from diffusion.utils.config import (
    ModelVideoConfig,
    SanaVideoConfig,
    TrainVideoConfig,
    VideoDataConfig,
    model_video_init_config,
)


@dataclass
class VideoDataCamCtrlConfig(VideoDataConfig):
    caption_proportion: Dict[str, Any] = field(default_factory=lambda: {"prompt": 1})
    hf_dataset_repo: Optional[str] = None
    hf_dataset_revision: Optional[str] = None
    hf_dataset_local_dir: str = "."
    hf_dataset_allow_patterns: Optional[List[str]] = None
    # Training-mixture fields; harmless at inference time but required so
    # pyrallis can deserialise the YAML.
    data_repeat: Optional[Dict[str, int]] = None
    shuffle_mode: str = "zip_group"
    external_caption_suffixes: List[str] = field(default_factory=list)
    return_raymap: bool = True
    use_plucker: bool = True
    vae_ratio: Tuple[int, int] = (4, 8)  # (time_downsample, spatial_downsample)
    cam_sample_strategy: str = "last"  # first, last
    use_relative_pose: bool = False
    normalize_poses: bool = False
    precompute_pos_embed: bool = False
    pos_embed_type: str = "wan_rope"
    attention_head_dim: int = 256
    patch_size: Tuple[int, int, int] = (1, 2, 2)
    max_seq_len: int = 1024
    pack_latents: bool = False
    precompute_cam_pos_embed: bool = False
    camctrl_type: str = "PRoPE"
    return_delta_actions: bool = False
    return_chunk_plucker: bool = False
    s3_config: Optional[Dict[str, Any]] = None
    s3_path_map: Optional[Dict[str, str]] = None


@dataclass
class ModelVideoCamCtrlConfig(ModelVideoConfig):
    camctrl_type: Optional[str] = None
    init_cam_from_base: bool = False
    cam_attn_compress: int = 1
    use_delta_actions: bool = False
    delta_action_dim: int = 64
    use_delta_translation: bool = False
    use_delta_pose_additive: bool = False
    delta_pose_additive_dim: int = 64
    use_chunk_plucker_input: bool = False
    use_chunk_plucker_post_attn: bool = False
    chunk_plucker_channels: int = 48
    chunk_plucker_post_attn_blocks: int = -1  # -1 = all blocks, N = first N blocks only
    fp32_norm: bool = False
    chunk_size: Optional[int] = None
    chunk_split_strategy: str = "uniform"
    conv_kernel_size: int = 4  # Temporal kernel size for ShortConvolution in GDN
    k_conv_only: bool = True  # Only apply ShortConvolution on K (skip Q and V)
    softmax_every_n: int = 4  # Hybrid GDN-Softmax: replace every N-th block with softmax (0=disabled)
    use_autograd_kernel: bool = False  # Switch Triton GDN blocks to autograd-enabled fused kernels (training mode)


@dataclass
class TrainVideoCamCtrlConfig(TrainVideoConfig):
    only_train_self_attn: bool = False
    only_train_cam_attn: bool = False
    # Per-sample mixture for chunk timestep sampling.
    chunk_mixture_probs: Optional[Dict[str, float]] = None
    mixed_finetune: bool = False
    main_lora_target_modules: Optional[List[str]] = None
    main_lora_include: List[str] = field(default_factory=lambda: [".attn.", ".mlp.", ".ffn."])
    main_lora_exclude: List[str] = field(default_factory=list)
    cam_branch_keywords: List[str] = field(
        default_factory=lambda: [
            "_proj_cam",
            "raymap_embedder",
            "delta_action_embedder",
            "delta_translation_embedder",
            "delta_pose_embedder",
            "delta_pose_proj",
            "plucker_embedder",
            "plucker_proj",
            "prope_proj",
            "rope_phase_",
        ]
    )
    max_steps: Optional[int] = None
    prefetch_factor: Optional[int] = None
    cam_branch_drop_prob: float = 0.0
    video_only_training_interval: int = 0
    train_batch_size_video_only: Optional[int] = None
    nocam_training_interval: int = 0
    train_batch_size_nocam: Optional[int] = None
    # Camera control validation settings
    camctrl_visualize: bool = False  # Enable camera control validation
    camctrl_val_data_path: str = "assets/camctrl_val_data.json"  # Path to validation data
    camctrl_val_cfg_scale: float = 6.0  # CFG scale for camera control validation
    camctrl_val_steps: int = 40  # Number of sampling steps for validation
    camctrl_val_wandb_scale: float = 1.5  # Upscale factor for wandb videos
    val_only: bool = False  # Run validation only and exit

    # Synchronize AR ``K`` (and the ``T_lat`` clamp) across data-parallel ranks
    # before sampling so every rank uses the same ``T_active = K + G``.
    ar_sync_K_across_ranks: bool = True


@dataclass
class SanaVideoCamCtrlConfig(SanaVideoConfig):
    data: VideoDataCamCtrlConfig
    model: ModelVideoCamCtrlConfig
    train: TrainVideoCamCtrlConfig
    video_only_data: Optional[VideoDataCamCtrlConfig] = None
    nocam_data: Optional[VideoDataCamCtrlConfig] = None
    tracker_project_name: str = "sana-video-camctrl"


def model_video_camctrl_init_config(config: SanaVideoCamCtrlConfig, latent_size: int = 32):
    return {
        "camctrl_type": config.model.camctrl_type,
        "cam_attn_compress": config.model.cam_attn_compress,
        "init_cam_from_base": config.model.init_cam_from_base,
        "use_delta_actions": config.model.use_delta_actions,
        "delta_action_dim": config.model.delta_action_dim,
        "use_delta_translation": config.model.use_delta_translation,
        "fp32_norm": config.model.fp32_norm,
        "chunk_size": config.model.chunk_size,
        "chunk_split_strategy": config.model.chunk_split_strategy,
        "conv_kernel_size": config.model.conv_kernel_size,
        "k_conv_only": config.model.k_conv_only,
        "use_delta_pose_additive": config.model.use_delta_pose_additive,
        "delta_pose_additive_dim": config.model.delta_pose_additive_dim,
        "use_chunk_plucker_input": config.model.use_chunk_plucker_input,
        "use_chunk_plucker_post_attn": config.model.use_chunk_plucker_post_attn,
        "chunk_plucker_channels": config.model.chunk_plucker_channels,
        "chunk_plucker_post_attn_blocks": config.model.chunk_plucker_post_attn_blocks,
        "softmax_every_n": config.model.softmax_every_n,
        "use_autograd_kernel": config.model.use_autograd_kernel,
        **model_video_init_config(config, latent_size=latent_size),
    }