import torch
import numpy as np
import os

from residual_diffusion import ResidualDiffusion
from transformer_denoiser import TransformerDenoiser


@torch.no_grad()
def generate_predictions(diffusion, model, sample, device, bounds, mean_r, std_r, num_preds=10, clip_len=750, clips_per_call=8, amp_dtype=torch.float16, guidance_weight=1.0):
    L = sample["length"]
    audio = sample["speaker_audio"] #(L, 768)
    dmm = sample["speaker_3dmm"] #(L, 58)
    emo = sample["speaker_emotion"] #(L, 25)

    if L % clip_len == 0:
        b = L // clip_len
        lengths = torch.tensor([clip_len] * b, dtype=torch.long)
    else:
        b = (L // clip_len) + 1
        pad = clip_len - (L % clip_len)
        audio = torch.cat((audio, torch.zeros(pad, audio.shape[-1])), dim=0)
        dmm = torch.cat((dmm, torch.zeros(pad, dmm.shape[-1])), dim=0)
        emo = torch.cat((emo, torch.zeros(pad, emo.shape[-1])), dim=0)
        lengths = torch.tensor([clip_len] * (b - 1) + [L % clip_len], dtype=torch.long)
    
    audio_c = audio.reshape(b, clip_len, -1).to(device)
    dmm_c = dmm.reshape(b, clip_len, -1).to(device)
    emo_c = emo.reshape(b, clip_len, -1).to(device)
    lengths = lengths.to(device)

    use_amp = (device == "cuda")
    per_clip = [None] * b

    for start in range(0, b, clips_per_call):
        stop = min(start + clips_per_call, b)
        g = stop - start

        audio_k = audio_c[start:stop].repeat(num_preds, 1, 1)
        dmm_k = dmm_c[start:stop].repeat(num_preds, 1, 1)
        emo_k = emo_c[start:stop].repeat(num_preds, 1, 1)
        lengths_k = lengths[start:stop].repeat(num_preds)

        with torch.amp.autocast("cuda", dtype=amp_dtype, enabled=use_amp):
            x = diffusion.ddim_sample(
                model=model, 
                shape=(num_preds*g, clip_len, 25), 
                speaker_audio=audio_k,
                speaker_3dmm=dmm_k,
                speaker_emotion=emo_k,
                mean_r=mean_r,
                std_r=std_r,
                motion_length=lengths_k,
                guidance_weight=guidance_weight,
            )
        
        x = x.float().reshape(num_preds, g, clip_len, 25)

        for i in range(g):
            per_clip[i + start] = x[:, i, :, :] #(num_preds, clip_len, 25)
        
    preds = torch.cat(per_clip, dim=1)[:, :L, :].cpu()

    preds[:, :, :15] = torch.round(preds[:, :, :15])
    preds = torch.clamp(
        preds, 
        min=torch.cat([torch.zeros(15), bounds["lower"]]), 
        max=torch.cat([torch.ones(15),  bounds["upper"]])
    )

    return preds

def load_data(modality, role, target_dim=None):
    path = os.path.join("sample_data", modality, role + ".npy")
    t = np.load(path).astype(np.float32)
    t = torch.from_numpy(t).squeeze()

    if target_dim is not None:
        assert t.dim() == target_dim

    return t

def main():
    if torch.cuda.is_available():
        device = "cuda"
        print("Connected to GPU")
    else:
        device = "cpu"
        print("Unable to connected to GPU; connected to CPU")
    
    diffusion = ResidualDiffusion()

    model = TransformerDenoiser().to(device).eval()
    ckpt = torch.load("./checkpoints/residual_diffusion_best.pt", map_location="cpu", weights_only=False)

    if isinstance(ckpt, dict) and "model" in ckpt:
        state_dict = ckpt["model"]
    elif isinstance(ckpt, dict) and "state_dict" in ckpt:
        state_dict = ckpt["state_dict"]
    else:
        state_dict = ckpt["model"]
    model.load_state_dict(state_dict)

    with open("./checkpoints/guidance_weight.txt", "r") as f:
        guidance_weight = f.readline()
    gw = float(guidance_weight.strip())

    bounds = torch.load("./checkpoints/bounds.pt")

    distribution = torch.load("./checkpoints/residual_distribution.pt")
    mean_r = distribution["mean"]
    std_r = distribution["std"]

    mean_face = np.load("./external_tools/mean_face.npy").astype(np.float32)
    std_face = np.load("./external_tools/std_face.npy").astype(np.float32)

    mean_face = torch.from_numpy(mean_face).unsqueeze(0)
    std_face = torch.from_numpy(std_face).unsqueeze(0)

    sp_audio = load_data("audio-features", "speaker", target_dim=2)
    sp_dmm = load_data("coefficients", "speaker", target_dim=2)
    sp_emo = load_data("facial-attributes", "speaker", target_dim=2)

    T = min(sp_audio.size(0), sp_dmm.size(0), sp_emo.size(0))
    sp_audio = sp_audio[:T]
    sp_dmm = sp_dmm[:T]
    sp_emo = sp_emo[:T]

    sp_dmm = (sp_dmm - mean_face) / std_face

    sample = {
        "length":T,
        "speaker_audio":sp_audio,
        "speaker_3dmm": sp_dmm,
        "speaker_emotion": sp_emo
    }

    print("Generating predictions\n\n")
    preds = generate_predictions(diffusion, model, sample, device, bounds, mean_r, std_r, guidance_weight=gw)

    print("Generation complete")
    print("Preds shape: ", preds.shape)
    np.save("results/generated_reactions.npy", preds.numpy())
    print("Saved successfully to results/generated_reactions.npy")

if __name__ == "__main__":
    main()