import pathlib
from vfi_utils import load_file_from_github_release, preprocess_frames, postprocess_frames, generic_frame_loop, InterpolationStateList
import typing
import torch
import torch.nn as nn
from comfy.model_management import soft_empty_cache, get_torch_device

MODEL_TYPE = pathlib.Path(__file__).parent.name
MODEL_FILE_NAMES = {
    "ssl": "eisai_ssl.pt",
    "dtm": "eisai_dtm.pt",
    "raft": "eisai_anime_interp_full.ckpt"
}

class EISAI(nn.Module):
    def __init__(self, model_file_names) -> None:
        from .eisai_arch import SoftsplatLite, DTM, RAFT
        super(EISAI, self).__init__()
        self.raft = RAFT(load_file_from_github_release(MODEL_TYPE, model_file_names["raft"]))
        self.raft.to(get_torch_device()).eval()

        self.ssl = SoftsplatLite()
        self.ssl.load_state_dict(torch.load(load_file_from_github_release(MODEL_TYPE, model_file_names["ssl"])))
        self.ssl.to(get_torch_device()).eval()

        self.dtm = DTM()
        self.dtm.load_state_dict(torch.load(load_file_from_github_release(MODEL_TYPE, model_file_names["dtm"])))
        self.dtm.to(get_torch_device()).eval()
    
    def forward(self, img0, img1, t):
        with torch.no_grad():
            flow0, _ = self.raft(img0, img1)
            flow1, _ = self.raft(img1, img0)
            x = {
                "images": torch.stack([img0, img1], dim=1),
                "flows": torch.stack([flow0, flow1], dim=1),
            }
            out_ssl, _ = self.ssl(x, t=t, return_more=True)
            out_dtm, _ = self.dtm(x, out_ssl, _, return_more=False)
        return out_dtm[:, :3]

class EISAI_VFI:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "ckpt_name": (["eisai"], ),
                "frames": ("IMAGE", ),
                "clear_cache_after_n_frames": ("INT", {"default": 10, "min": 1, "max": 1000}),
                "multiplier": ("INT", {"default": 2, "min": 2, "max": 1000}),
            },
            "optional": {
                "optional_interpolation_states": ("INTERPOLATION_STATES", )
            }
        }

    RETURN_TYPES = ("IMAGE", )
    FUNCTION = "vfi"
    CATEGORY = "ComfyUI-Frame-Interpolation/VFI"

    def vfi(
        self,
        ckpt_name: typing.AnyStr,
        frames: torch.Tensor,
        clear_cache_after_n_frames = 10,
        multiplier: typing.SupportsInt = 2,
        optional_interpolation_states: InterpolationStateList = None,
        **kwargs
    ):
        interpolation_model = EISAI(MODEL_FILE_NAMES)
        interpolation_model.eval().to(get_torch_device())
        frames = preprocess_frames(frames)
        
        def return_middle_frame(frame_0, frame_1, timestep, model):
            return model(frame_0, frame_1, t=timestep)
        
        scale = 1
        
        args = [interpolation_model, scale]
        out = postprocess_frames(
            generic_frame_loop(type(self).__name__, frames, clear_cache_after_n_frames, multiplier, return_middle_frame, *args, 
                               interpolation_states=optional_interpolation_states, dtype=torch.float32)
        )
        return (out,)