import lpips
import torch
import torch.optim as optim
from pytorch_msssim import ms_ssim
from torchcfm.conditional_flow_matching import (
    ConditionalFlowMatcher,
    ExactOptimalTransportConditionalFlowMatcher,
    TargetConditionalFlowMatcher,
    VariancePreservingConditionalFlowMatcher,
)
from torchdiffeq import odeint

import wandb
from src.dataset.preprocess import denormalize
from src.gan import PatchDiscriminator, gan_disc_loss

lecam_loss_weight = 0.1
lecam_anchor_real_logits = 0.0
lecam_anchor_fake_logits = 0.0
lecam_beta = 0.9
use_lecam = True


class RF:
    def __init__(self, model, fm="otcfm", loss="mse"):
        self.model = model
        self.loss = loss
        self.iter = 0

        sigma = 0.0
        if fm == "otcfm":
            self.fm = ExactOptimalTransportConditionalFlowMatcher(sigma=sigma)
        elif fm == "icfm":
            self.fm = ConditionalFlowMatcher(sigma=sigma)
        elif fm == "fm":
            self.fm = TargetConditionalFlowMatcher(sigma=sigma)
        elif fm == "si":
            self.fm = VariancePreservingConditionalFlowMatcher(sigma=sigma)
        else:
            raise NotImplementedError(
                f"Unknown model {fm}, must be one of ['otcfm', 'icfm', 'fm', 'si']"
            )

        self.lpips = lpips.LPIPS(net="vgg").to("cuda:1")
        self.lpips2 = lpips.LPIPS(net="alex").to("cuda:1")

        discriminator = PatchDiscriminator().to("cuda:1")
        discriminator.requires_grad_(True)
        self.discriminator = discriminator
        self.optimizer_D = optim.AdamW(
            discriminator.parameters(),
            lr=2e-4,
            weight_decay=1e-3,
            betas=(0.9, 0.95),
        )

    def gan_loss(self, real, fake):
        global lecam_beta, lecam_anchor_real_logits, lecam_anchor_fake_logits, use_lecam

        real_preds = self.discriminator(real)
        fake_preds = self.discriminator(fake.detach())
        d_loss, avg_real_logits, avg_fake_logits, disc_acc = gan_disc_loss(
            real_preds, fake_preds, "hinge"
        )

        lecam_anchor_real_logits = (
            lecam_beta * lecam_anchor_real_logits + (1 - lecam_beta) * avg_real_logits
        )
        lecam_anchor_fake_logits = (
            lecam_beta * lecam_anchor_fake_logits + (1 - lecam_beta) * avg_fake_logits
        )
        total_d_loss = d_loss.mean()
        d_loss_item = total_d_loss.item()
        if use_lecam:
            # penalize the real logits to fake and fake logits to real.
            lecam_loss = (real_preds - lecam_anchor_fake_logits).pow(2).mean() + (
                fake_preds - lecam_anchor_real_logits
            ).pow(2).mean()
            lecam_loss_item = lecam_loss.item()
            total_d_loss = total_d_loss + lecam_loss * lecam_loss_weight

            wandb.log(
                {
                    "gan/lecam_loss": lecam_loss_item,
                    "gan/lecam_anchor_real_logits": lecam_anchor_real_logits,
                    "gan/lecam_anchor_fake_logits": lecam_anchor_fake_logits,
                }
            )

        wandb.log(
            {
                "gan/discriminator_loss": d_loss_item,
                "gan/discriminator_accuracy": disc_acc,
            }
        )

        self.optimizer_D.zero_grad()
        total_d_loss.backward(retain_graph=True)
        self.optimizer_D.step()

    def forward(self, gt, cloud, condition=False):
        t, xt, ut = self.fm.sample_location_and_conditional_flow(  # type: ignore
            cloud if not condition else torch.randn_like(cloud), gt
        )
        if condition:
            vt = self.model(xt, t, cloud)
        else:
            vt, _ = self.model(xt, t)

        if self.loss == "mse":
            loss = ((vt - ut) ** 2).mean(dim=list(range(1, len(gt.shape))))
            loss_list = {"train/mse": loss.mean().item()}
        elif self.loss == "lpips_mse":
            mse = ((vt - ut) ** 2).mean(dim=list(range(1, len(gt.shape))))
            lpips = self.lpips(
                denormalize(gt) * 2 - 1,
                denormalize(xt + (1 - t[:, None, None, None]) * vt) * 2 - 1,
            )
            loss_list = {
                "train/mse": mse.mean().item(),
                "train/lpips": lpips.mean().item(),
            }
            loss = mse + lpips * 2.0
        elif self.loss == "gan_lpips_mse":
            self.iter += 1
            # if self.iter % 4 == 0:
            #     self.gan_loss(
            #         denormalize(gt),
            #         denormalize(xt + (1 - t[:, None, None, None]) * vt),
            #     )
            mse = ((vt - ut) ** 2).mean(dim=list(range(1, len(gt.shape))))
            lpips = self.lpips(
                denormalize(gt) * 2 - 1,
                denormalize(xt + (1 - t[:, None, None, None]) * vt) * 2 - 1,
            )
            alexlpips = self.lpips2(
                denormalize(gt) * 2 - 1,
                denormalize(xt + (1 - t[:, None, None, None]) * vt) * 2 - 1,
            )
            # gan = -self.discriminator(
            #     denormalize(xt + (1 - t[:, None, None, None]) * vt),
            # ).mean(-1)
            ssim = 1 - ms_ssim(
                denormalize(gt),
                denormalize(xt + (1 - t[:, None, None, None]) * vt),
                data_range=1.0,
                size_average=False,
            )
            loss_list = {
                "train/mse": mse.mean().item(),
                "train/lpips": lpips.mean().item(),
                "train/alexlpips": alexlpips.mean().item(),
                # "train/gan": gan.mean().item(),
                "train/ssim": ssim.mean().item(),
            }
            loss = mse + lpips * 2.0 + alexlpips + ssim
        else:
            raise Exception

        tlist = loss.detach().cpu().reshape(-1).tolist()
        ttloss = [(tv, tloss) for tv, tloss in zip(t, tlist)]

        return loss.mean(), ttloss, loss_list

    @torch.no_grad()
    def sample(
        self, cloud, tol=1e-5, integration="dopri5", condition=False
    ) -> torch.Tensor:
        t_span = torch.linspace(0, 1, 2, device=cloud.device)
        if condition:
            x = torch.randn_like(cloud)
            traj = odeint(
                lambda t, x: self.model(x, t, cloud),
                x,
                t_span,
                rtol=tol,
                atol=tol,
                method=integration,
            )
        else:
            traj = odeint(
                lambda t, x: self.model(x, t)[0],
                cloud,
                t_span,
                rtol=tol,
                atol=tol,
                method=integration,
            )

        return [traj[i] for i in range(traj.shape[0])]  # type: ignore
