import torch
from enum import Enum
import gc
import numpy as np
import jax.numpy as jnp
import jax

from PIL import Image
from typing import List

from flax.training.common_utils import shard
from flax.jax_utils import replicate
from flax import jax_utils
import einops

from transformers import CLIPTokenizer, CLIPFeatureExtractor, FlaxCLIPTextModel
from diffusers import (
    FlaxDDIMScheduler,
    FlaxAutoencoderKL,
    FlaxStableDiffusionControlNetPipeline,
    StableDiffusionPipeline,
    FlaxUNet2DConditionModel as VanillaFlaxUNet2DConditionModel,
)
from text_to_animation.models.unet_2d_condition_flax import (
    FlaxUNet2DConditionModel
)
from diffusers import FlaxControlNetModel

from text_to_animation.pipelines.text_to_video_pipeline_flax import (
    FlaxTextToVideoPipeline,
)

import utils.utils as utils
import utils.gradio_utils as gradio_utils
import os

on_huggingspace = os.environ.get("SPACE_AUTHOR_NAME") == "PAIR"

unshard = lambda x: einops.rearrange(x, "d b ... -> (d b) ...")


class ModelType(Enum):
    Text2Video = 1
    ControlNetPose = 2
    StableDiffusion = 3


def replicate_devices(array):
    return jnp.expand_dims(array, 0).repeat(jax.device_count(), 0)


class ControlAnimationModel:
    def __init__(self, dtype, **kwargs):
        self.dtype = dtype
        self.rng = jax.random.PRNGKey(0)
        self.pipe = None
        self.model_type = None

        self.states = {}
        self.model_name = ""

    def set_model(
        self,
        model_id: str,
        **kwargs,
    ):
        if hasattr(self, "pipe") and self.pipe is not None:
            del self.pipe
            self.pipe = None
        gc.collect()

        controlnet, controlnet_params = FlaxControlNetModel.from_pretrained(
            "fusing/stable-diffusion-v1-5-controlnet-openpose",
            from_pt=True,
            dtype=jnp.float16,
        )

        scheduler, scheduler_state = FlaxDDIMScheduler.from_pretrained(
            model_id, subfolder="scheduler", from_pt=True
        )
        tokenizer = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer")
        feature_extractor = CLIPFeatureExtractor.from_pretrained(
            model_id, subfolder="feature_extractor"
        )
        unet, unet_params = FlaxUNet2DConditionModel.from_pretrained(
            model_id, subfolder="unet", from_pt=True, dtype=self.dtype
        )
        unet_vanilla = VanillaFlaxUNet2DConditionModel.from_config(
            model_id, subfolder="unet", from_pt=True, dtype=self.dtype
        )
        vae, vae_params = FlaxAutoencoderKL.from_pretrained(
            model_id, subfolder="vae", from_pt=True, dtype=self.dtype
        )
        text_encoder = FlaxCLIPTextModel.from_pretrained(
            model_id, subfolder="text_encoder", from_pt=True, dtype=self.dtype
        )
        self.pipe = FlaxTextToVideoPipeline(
            vae=vae,
            text_encoder=text_encoder,
            tokenizer=tokenizer,
            unet=unet,
            unet_vanilla=unet_vanilla,
            controlnet=controlnet,
            scheduler=scheduler,
            safety_checker=None,
            feature_extractor=feature_extractor,
        )
        self.params = {
            "unet": unet_params,
            "vae": vae_params,
            "scheduler": scheduler_state,
            "controlnet": controlnet_params,
            "text_encoder": text_encoder.params,
        }
        self.p_params = jax_utils.replicate(self.params)
        self.model_name = model_id

    def generate_initial_frames(
        self,
        prompt: str,
        video_path: str,
        n_prompt: str = "",
        num_imgs: int = 4,
        resolution: int = 512,
        model_id: str = "runwayml/stable-diffusion-v1-5",
    ) -> List[Image.Image]:
        self.set_model(model_id=model_id)

        video_path = gradio_utils.motion_to_video_path(video_path)

        added_prompt = "high quality, best quality, HD, clay stop-motion, claymation, HQ, masterpiece, art, smooth"
        prompts = added_prompt + ", " + prompt

        added_n_prompt = "longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer difits, cropped, worst quality, low quality, deformed body, bloated, ugly"
        negative_prompts = added_n_prompt + ", " + n_prompt

        video, fps = utils.prepare_video(
            video_path, resolution, None, self.dtype, False, output_fps=4
        )
        control = utils.pre_process_pose(video, apply_pose_detect=False)

        seeds = [seed for seed in jax.random.randint(self.rng, [num_imgs], 0, 65536)]
        prngs = [jax.random.PRNGKey(seed) for seed in seeds]
        print(seeds)
        images = self.pipe.generate_starting_frames(
            params=self.p_params,
            prngs=prngs,
            controlnet_image=control,
            prompt=prompts,
            neg_prompt=negative_prompts,
        )

        images = [np.array(images[i]) for i in range(images.shape[0])]

        return images

    def generate_video_from_frame(self, controlnet_video, prompt, seed, neg_prompt=""):
        # generate a video using the seed provided
        prng_seed = jax.random.PRNGKey(seed)
        len_vid = controlnet_video.shape[0]
        # print(f"Generating video from prompt {'<aardman> style '+ prompt}, with {controlnet_video.shape[0]} frames and prng seed {seed}")
        added_prompt = "high quality, best quality, HD, clay stop-motion, claymation, HQ, masterpiece, art, smooth"
        prompts = added_prompt + ", " + prompt

        added_n_prompt = "longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer difits, cropped, worst quality, low quality, deformed body, bloated, ugly"
        negative_prompts = added_n_prompt + ", " + neg_prompt
        
        # prompt_ids = self.pipe.prepare_text_inputs(["aardman style "+ prompt]*len_vid)
        # n_prompt_ids = self.pipe.prepare_text_inputs([neg_prompt]*len_vid)
        
        prompt_ids = self.pipe.prepare_text_inputs([prompts]*len_vid)
        n_prompt_ids = self.pipe.prepare_text_inputs([negative_prompts]*len_vid)
        prng = replicate_devices(prng_seed) #jax.random.split(prng, jax.device_count())
        image = replicate_devices(controlnet_video)
        prompt_ids = replicate_devices(prompt_ids)
        n_prompt_ids = replicate_devices(n_prompt_ids)
        motion_field_strength_x = replicate_devices(jnp.array(3))
        motion_field_strength_y = replicate_devices(jnp.array(4))
        smooth_bg_strength = replicate_devices(jnp.array(0.8))
        vid = (self.pipe(image=image,
                        prompt_ids=prompt_ids,
                        neg_prompt_ids=n_prompt_ids, 
                        params=self.p_params,
                        prng_seed=prng,
                        jit = True,
                        smooth_bg_strength=smooth_bg_strength,
                        motion_field_strength_x=motion_field_strength_x,
                        motion_field_strength_y=motion_field_strength_y,
                        ).images)[0]
        return utils.create_gif(np.array(vid), 4, path=None, watermark=None)


    def generate_animation(
        self,
        prompt, #: str,
        initial_frame_index, #: int,
        input_video_path, #: str,
        model_link = None,#: str = "dreamlike-art/dreamlike-photoreal-2.0",
        motion_field_strength_x = 12,#: int = 12,
        motion_field_strength_y= 12,#: int = 12,
        t0= 44,#: int = 44,
        t1= 47,#: int = 47,
        n_prompt= "",#: str = "",
        chunk_size= 8, #: int = 8,
        video_length = 8, #: int = 8,
        merging_ratio = 0., #: float = 0.0,
        seed= 0,#: int = 0,
        resolution=512,#: int = 512,
        fps=2,#: int = 2,
        use_cf_attn=True,#: bool = True,
        use_motion_field=True,#: bool = True,
        smooth_bg=False,#: bool = False,
        smooth_bg_strength=0.4,#: float = 0.4,
        path=None,#: str = None,
    ):
        video_path = gradio_utils.motion_to_video_path(input_video_path)

        # added_prompt = 'best quality, HD, clay stop-motion, claymation, HQ, masterpiece, art, smooth'
        # added_prompt = 'high quality, anatomically correct, clay stop-motion, aardman, claymation, smooth'
        added_n_prompt = "longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer difits, cropped, worst quality, low quality, deformed body, bloated, ugly"
        negative_prompts = added_n_prompt + ", " + n_prompt

        video, fps = utils.prepare_video(
            video_path, resolution, None, self.dtype, False, output_fps=4
        )
        control = utils.pre_process_pose(video, apply_pose_detect=False)
        len_vid, _, h, w = video.shape
        prng_seed = jax.random.PRNGKey(seed)
        prompts = prompt
        prompt_ids = self.pipe.prepare_text_inputs([prompts]*len_vid)
        n_prompt_ids = self.pipe.prepare_text_inputs([negative_prompts]*len_vid)
        prng = replicate_devices(prng_seed) #jax.random.split(prng, jax.device_count())
        image = replicate_devices(control)
        prompt_ids = replicate_devices(prompt_ids)
        n_prompt_ids = replicate_devices(n_prompt_ids)
        motion_field_strength_x = replicate_devices(jnp.array(motion_field_strength_x))
        motion_field_strength_y = replicate_devices(jnp.array(motion_field_strength_y))
        smooth_bg_strength = replicate_devices(jnp.array(smooth_bg_strength))
        vid = (self.pipe(image=image,
                        prompt_ids=prompt_ids,
                        neg_prompt_ids=n_prompt_ids, 
                        params=self.p_params,
                        prng_seed=prng,
                        jit = True,
                        smooth_bg_strength=smooth_bg_strength,
                        motion_field_strength_x=motion_field_strength_x,
                        motion_field_strength_y=motion_field_strength_y,
                        ).images)[0]
        return utils.create_gif(np.array(vid), 4, path=None, watermark=None)