from typing import Union, Any, Dict, List, Optional, Callable
from . import pl_module_extension
from .diffusers_conditional.models.controlnet.image_embedder import AbstractEncoder
from .requires_grad_setter import LayerConfig as LayerConfigNew
from . import video_noise_generator


def auto_str(cls):
    def __str__(self):
        return "%s(%s)" % (
            type(self).__name__,
            ", ".join("%s=%s" % item for item in vars(self).items()),
        )

    cls.__str__ = __str__
    return cls


class LayerConfig:
    def __init__(
        self,
        update_with_full_lr: Optional[Union[List[str], List[List[str]]]] = None,
        exclude: Optional[List[str]] = None,
        deactivate_all_grads: bool = True,
    ) -> None:
        self.deactivate_all_grads = deactivate_all_grads
        if exclude is not None:
            self.exclude = exclude
        if update_with_full_lr is not None:
            self.update_with_full_lr = update_with_full_lr

    def __str__(self) -> str:
        str = f"Deactivate all gradients first={self.deactivate_all_grads}. "
        if hasattr(self, "update_with_full_lr"):
            str += f"Then activating gradients for: {self.update_with_full_lr}. "
        if hasattr(self, "exclude"):
            str += f"Finally, excluding: {self.exclude}. "
        return str


class OptimizerParams:
    def __init__(
        self,
        learning_rate: float,
        # Default value due to legacy
        layers_config: Union[LayerConfig, LayerConfigNew] = None,
        layers_config_base: LayerConfig = None,  # Default value due to legacy
        use_warmup: bool = False,
        warmup_steps: int = 10000,
        warmup_start_factor: float = 1e-5,
        learning_rate_spatial: float = 0.0,
        use_8_bit_adam: bool = False,
        noise_generator: Union[
            pl_module_extension.NoiseGenerator, video_noise_generator.NoiseGenerator
        ] = None,
        noise_decomposition: pl_module_extension.NoiseDecomposition = None,
        perceptual_loss: bool = False,
        noise_offset: float = 0.0,
        split_opt_by_node: bool = False,
        reset_prediction_type_to_eps: bool = False,
        train_val_sampler_may_differ: bool = False,
        measure_similarity: bool = False,
        similarity_loss: bool = False,
        similarity_loss_weight: float = 1.0,
        loss_conditional_weight: float = 0.0,
        loss_conditional_weight_convex: bool = False,
        loss_conditional_change_after_step: int = 0,
        mask_conditional_frames: bool = False,
        sample_from_noise: bool = True,
        mask_alternating: bool = False,
        uncondition_freq: int = -1,
        no_text_condition_control: bool = False,
        inject_image_into_input: bool = False,
        inject_at_T: bool = False,
        resampling_steps: int = 1,
        control_freq_in_resample: int = 1,
        resample_to_T: bool = False,
        adaptive_loss_reweight: bool = False,
        load_resampler_from_ckpt: str = "",
        skip_controlnet_branch: bool = False,
        use_fps_conditioning: bool = False,
        num_frame_embeddings_range: int = 16,
        start_frame_training: int = 0,
        start_frame_ctrl: int = 0,
        load_trained_base_model_and_resampler_from_ckpt: str = "",
        load_trained_controlnet_from_ckpt: str = "",
        # fill_up_frame_to_video: bool = False,
    ) -> None:
        self.use_warmup = use_warmup
        self.warmup_steps = warmup_steps
        self.warmup_start_factor = warmup_start_factor
        self.learning_rate_spatial = learning_rate_spatial
        self.learning_rate = learning_rate
        self.use_8_bit_adam = use_8_bit_adam
        self.layers_config = layers_config
        self.noise_generator = noise_generator
        self.perceptual_loss = perceptual_loss
        self.noise_decomposition = noise_decomposition
        self.noise_offset = noise_offset
        self.split_opt_by_node = split_opt_by_node
        self.reset_prediction_type_to_eps = reset_prediction_type_to_eps
        self.train_val_sampler_may_differ = train_val_sampler_may_differ
        self.measure_similarity = measure_similarity
        self.similarity_loss = similarity_loss
        self.similarity_loss_weight = similarity_loss_weight
        self.loss_conditional_weight = loss_conditional_weight
        self.loss_conditional_change_after_step = loss_conditional_change_after_step
        self.mask_conditional_frames = mask_conditional_frames
        self.loss_conditional_weight_convex = loss_conditional_weight_convex
        self.sample_from_noise = sample_from_noise
        self.layers_config_base = layers_config_base
        self.mask_alternating = mask_alternating
        self.uncondition_freq = uncondition_freq
        self.no_text_condition_control = no_text_condition_control
        self.inject_image_into_input = inject_image_into_input
        self.inject_at_T = inject_at_T
        self.resampling_steps = resampling_steps
        self.control_freq_in_resample = control_freq_in_resample
        self.resample_to_T = resample_to_T
        self.adaptive_loss_reweight = adaptive_loss_reweight
        self.load_resampler_from_ckpt = load_resampler_from_ckpt
        self.skip_controlnet_branch = skip_controlnet_branch
        self.use_fps_conditioning = use_fps_conditioning
        self.num_frame_embeddings_range = num_frame_embeddings_range
        self.start_frame_training = start_frame_training
        self.load_trained_base_model_and_resampler_from_ckpt = (
            load_trained_base_model_and_resampler_from_ckpt
        )
        self.load_trained_controlnet_from_ckpt = load_trained_controlnet_from_ckpt
        self.start_frame_ctrl = start_frame_ctrl
        if start_frame_ctrl < 0:
            print("new format start frame cannot be negative")
            exit()

        # self.fill_up_frame_to_video = fill_up_frame_to_video

    @property
    def learning_rate_spatial(self):
        return self._learning_rate_spatial

    # legacy code that maps the state None or '-1' to '0.0'
    # so 0.0 indicated no spatial learning rate is selected
    @learning_rate_spatial.setter
    def learning_rate_spatial(self, value):
        if value is None or value == -1:
            value = 0
        self._learning_rate_spatial = value


# Legacy class
class SchedulerParams:
    def __init__(
        self,
        use_warmup: bool = False,
        warmup_steps: int = 10000,
        warmup_start_factor: float = 1e-5,
    ) -> None:
        self.use_warmup = use_warmup
        self.warmup_steps = warmup_steps
        self.warmup_start_factor = warmup_start_factor


class CrossFrameAttentionParams:

    def __init__(self, attent_on: List[int], masking=False) -> None:
        self.attent_on = attent_on
        self.masking = masking


class InferenceParams:
    def __init__(
        self,
        width: int,
        height: int,
        video_length: int,
        guidance_scale: float = 7.5,
        use_dec_scaling: bool = True,
        frame_rate: int = 2,
        num_inference_steps: int = 50,
        eta: float = 0.0,
        n_autoregressive_generations: int = 1,
        mode: str = "long_video",
        start_from_real_input: bool = True,
        eval_loss_metrics: bool = False,
        scheduler_cls: str = "",
        negative_prompt: str = "",
        conditioning_from_all_past: bool = False,
        validation_samples: int = 80,
        conditioning_type: str = "last_chunk",
        result_formats: List[str] = ["eval_gif", "gif", "mp4"],
        concat_video: bool = True,
        seed: int = 33,
    ):
        self.width = width
        self.height = height
        self.video_length = (
            video_length if isinstance(video_length, int) else int(video_length)
        )
        self.guidance_scale = guidance_scale
        self.use_dec_scaling = use_dec_scaling
        self.frame_rate = frame_rate
        self.num_inference_steps = num_inference_steps
        self.eta = eta
        self.negative_prompt = negative_prompt
        self.n_autoregressive_generations = n_autoregressive_generations
        self.mode = mode
        self.start_from_real_input = start_from_real_input
        self.eval_loss_metrics = eval_loss_metrics
        self.scheduler_cls = scheduler_cls
        self.conditioning_from_all_past = conditioning_from_all_past
        self.validation_samples = validation_samples
        self.conditioning_type = conditioning_type
        self.result_formats = result_formats
        self.concat_video = concat_video
        self.seed = seed

    def to_dict(self):

        keys = [
            entry
            for entry in dir(self)
            if not callable(getattr(self, entry)) and not entry.startswith("__")
        ]

        result_dict = {}
        for key in keys:
            result_dict[key] = getattr(self, key)
        return result_dict


@auto_str
class AttentionMaskParams:

    def __init__(
        self,
        temporal_self_attention_only_on_conditioning: bool = False,
        temporal_self_attention_mask_included_itself: bool = False,
        spatial_attend_on_condition_frames: bool = False,
        temp_attend_on_neighborhood_of_condition_frames: bool = False,
        temp_attend_on_uncond_include_past: bool = False,
    ) -> None:
        self.temporal_self_attention_mask_included_itself = (
            temporal_self_attention_mask_included_itself
        )
        self.spatial_attend_on_condition_frames = spatial_attend_on_condition_frames
        self.temp_attend_on_neighborhood_of_condition_frames = (
            temp_attend_on_neighborhood_of_condition_frames
        )
        self.temporal_self_attention_only_on_conditioning = (
            temporal_self_attention_only_on_conditioning
        )
        self.temp_attend_on_uncond_include_past = temp_attend_on_uncond_include_past

        assert (
            not temp_attend_on_neighborhood_of_condition_frames
            or not temporal_self_attention_only_on_conditioning
        )


class UNetParams:

    def __init__(
        self,
        conditioning_embedding_out_channels: List[int],
        ckpt_spatial_layers: str = "",
        pipeline_repo: str = "",
        unet_from_diffusers: bool = True,
        spatial_latent_input: bool = False,
        num_frame_conditioning: int = 1,
        pipeline_class: str = "t2v_enhanced.model.model.controlnet.pipeline_text_to_video_w_controlnet_synth.TextToVideoSDPipeline",
        frame_expansion: str = "last_frame",
        downsample_controlnet_cond: bool = True,
        num_frames: int = 1,
        pre_transformer_in_cond: bool = False,
        num_tranformers: int = 1,
        zero_conv_3d: bool = False,
        merging_mode: str = "addition",
        compute_only_conditioned_frames: bool = False,
        condition_encoder: str = "",
        zero_conv_mode: str = "2d",
        clean_model: bool = False,
        merging_mode_base: str = "addition",
        attention_mask_params: AttentionMaskParams = None,
        attention_mask_params_base: AttentionMaskParams = None,
        modelscope_input_format: bool = True,
        temporal_self_attention_only_on_conditioning: bool = False,
        temporal_self_attention_mask_included_itself: bool = False,
        use_post_merger_zero_conv: bool = False,
        weight_control_sample: float = 1.0,
        use_controlnet_mask: bool = False,
        random_mask_shift: bool = False,
        random_mask: bool = False,
        use_resampler: bool = False,
        unet_from_pipe: bool = False,
        unet_operates_on_2d: bool = False,
        image_encoder: str = "CLIP",
        use_standard_attention_processor: bool = True,
        num_frames_before_chunk: int = 0,
        resampler_type: str = "single_frame",
        resampler_cls: str = "",
        resampler_merging_layers: int = 1,
        image_encoder_obj: AbstractEncoder = None,
        cfg_text_image: bool = False,
        aggregation: str = "last_out",
        resampler_random_shift: bool = False,
        img_cond_alpha_per_frame: bool = False,
        num_control_input_frames: int = -1,
        use_image_encoder_normalization: bool = False,
        use_of: bool = False,
        ema_param: float = -1.0,
        concat: bool = False,
        use_image_tokens_main: bool = True,
        use_image_tokens_ctrl: bool = False,
    ):

        self.ckpt_spatial_layers = ckpt_spatial_layers
        self.pipeline_repo = pipeline_repo
        self.unet_from_diffusers = unet_from_diffusers
        self.spatial_latent_input = spatial_latent_input
        self.pipeline_class = pipeline_class
        self.num_frame_conditioning = num_frame_conditioning
        if num_control_input_frames == -1:
            self.num_control_input_frames = num_frame_conditioning
        else:
            self.num_control_input_frames = num_control_input_frames

        self.conditioning_embedding_out_channels = conditioning_embedding_out_channels
        self.frame_expansion = frame_expansion
        self.downsample_controlnet_cond = downsample_controlnet_cond
        self.num_frames = num_frames
        self.pre_transformer_in_cond = pre_transformer_in_cond
        self.num_tranformers = num_tranformers
        self.zero_conv_3d = zero_conv_3d
        self.merging_mode = merging_mode
        self.compute_only_conditioned_frames = compute_only_conditioned_frames
        self.clean_model = clean_model
        self.condition_encoder = condition_encoder
        self.zero_conv_mode = zero_conv_mode
        self.merging_mode_base = merging_mode_base
        self.modelscope_input_format = modelscope_input_format
        assert (
            not temporal_self_attention_only_on_conditioning
        ), "This parameter is only here for backward compatibility. Set AttentionMaskParams instead."
        assert (
            not temporal_self_attention_mask_included_itself
        ), "This parameter is only here for backward compatibility. Set AttentionMaskParams instead."
        if attention_mask_params is not None and attention_mask_params_base is None:
            attention_mask_params_base = attention_mask_params
        if attention_mask_params is None:
            attention_mask_params = AttentionMaskParams()
        if attention_mask_params_base is None:
            attention_mask_params_base = AttentionMaskParams()
        self.attention_mask_params = attention_mask_params
        self.attention_mask_params_base = attention_mask_params_base
        self.weight_control_sample = weight_control_sample
        self.use_controlnet_mask = use_controlnet_mask
        self.random_mask_shift = random_mask_shift
        self.random_mask = random_mask
        self.use_resampler = use_resampler
        self.unet_from_pipe = unet_from_pipe
        self.unet_operates_on_2d = unet_operates_on_2d
        self.image_encoder = image_encoder_obj
        self.use_standard_attention_processor = use_standard_attention_processor
        self.num_frames_before_chunk = num_frames_before_chunk
        self.resampler_type = resampler_type
        self.resampler_cls = resampler_cls
        self.resampler_merging_layers = resampler_merging_layers
        self.cfg_text_image = cfg_text_image
        self.aggregation = aggregation
        self.resampler_random_shift = resampler_random_shift
        self.img_cond_alpha_per_frame = img_cond_alpha_per_frame
        self.use_image_encoder_normalization = use_image_encoder_normalization
        self.use_of = use_of
        self.ema_param = ema_param
        self.concat = concat
        self.use_image_tokens_main = use_image_tokens_main
        self.use_image_tokens_ctrl = use_image_tokens_ctrl
        assert not use_post_merger_zero_conv

        if spatial_latent_input:
            assert (
                unet_from_diffusers
            ), "Spatial latent input only implemented by original diffusers model. Set 'model.unet_params.unet_from_diffusers=True'."
