import os
import time
from copy import deepcopy
from types import SimpleNamespace

import torch
import torch.optim as optim
from torch import amp
from torchvision.utils import save_image

from utils import AvgMeter, EMA, load_config, psnr, resolve_device, set_seed, ssim
from dataloader import get_dataloaders, VARIANT_TO_IDX
from model import build_model


def train(config_path: str = "config.yaml") -> None:
    cfg = load_config(config_path)
    data_cfg = cfg.data
    training_cfg = cfg.training
    eval_cfg = getattr(cfg, "evaluation", SimpleNamespace(save_dir="eval_samples"))
    model_cfg = cfg.model

    out_dir = training_cfg.out_dir
    os.makedirs(out_dir, exist_ok=True)

    set_seed(training_cfg.seed)
    device = resolve_device(training_cfg.device)
    use_cuda = isinstance(device, str) and device.startswith("cuda")

    attention_cfg = getattr(model_cfg, "attention", SimpleNamespace())
    print(
        f"[Config] Using device={device}, image_size={data_cfg.image_size}, "
        f"backbone={model_cfg.backbone}, attention={getattr(attention_cfg, 'enabled', False)}"
    )

    test_workers = getattr(eval_cfg, "num_workers", getattr(data_cfg, "test_num_workers", 2))
    train_loader, test_loader = get_dataloaders(
        data_root=data_cfg.root,
        batch_size=training_cfg.batch_size,
        image_size=data_cfg.image_size,
        train_variant=data_cfg.train_variant,
        test_variant=data_cfg.test_variant,
        train_num_workers=data_cfg.train_num_workers,
        test_num_workers=test_workers,
        pin_memory=data_cfg.pin_memory,
    )
    
    # Build model from configuration
    model = build_model(cfg, device=device)

    ema_cfg = getattr(model_cfg, "ema", SimpleNamespace(enabled=True, decay=0.999))
    use_ema = getattr(ema_cfg, "enabled", True)
    ema = EMA(model, decay=getattr(ema_cfg, "decay", 0.999)) if use_ema else None

    opt = optim.AdamW(model.parameters(), lr=training_cfg.learning_rate, betas=(0.9, 0.999))
    scaler = amp.GradScaler("cuda", enabled=use_cuda)

    model.train()

    total_start = time.time()
    best_rec = float("inf")
    best_ckpt_path = os.path.join(out_dir, "model_best.pt")

    preview_interval = getattr(training_cfg, "preview_eval_interval", 10)
    preview_samples = getattr(training_cfg, "preview_eval_samples", 1)
    preview_save_dir = getattr(training_cfg, "preview_save_dir", getattr(eval_cfg, "save_dir", "preview_samples"))
    preview_save_images = bool(getattr(training_cfg, "preview_save_images", True))

    def _to_variant_idx(batch_dict: dict) -> torch.Tensor | None:
        variant = batch_dict.get("variant")
        if variant is None:
            return None
        if torch.is_tensor(variant):
            return variant.to(device, dtype=torch.long)
        return torch.tensor([VARIANT_TO_IDX[v] for v in variant], device=device, dtype=torch.long)

    for epoch in range(1, training_cfg.epochs + 1):
        epoch_start = time.time()
        loss_meter = AvgMeter()
        rec_meter = AvgMeter()

        for batch in train_loader:
            y = batch["y"].to(device)
            x = batch["gt"].to(device)
            B = x.size(0)
            variant_idx = _to_variant_idx(batch)

            t = torch.randint(0, model.cfg.T, (B,), device=device, dtype=torch.long)
            x_t, noise = model.q_sample(x, t)

            with amp.autocast("cuda", enabled=use_cuda):
                eps_pred = model.predict_eps(x_t, t, y, variant_idx)
                loss_noise = ((noise - eps_pred) ** 2).mean()

                a_bar = model.alphas_cumprod[t].view(-1, 1, 1, 1)
                x0_hat = (x_t - torch.sqrt(1.0 - a_bar) * eps_pred) / torch.sqrt(a_bar)
                x0_hat = torch.clamp(x0_hat, 0.0, 1.0)
                loss_rec = (x0_hat - x).abs().mean()

                if training_cfg.lambda_cons > 0.0 and B >= 2:
                    x0_a, x0_b = x0_hat[::2], x0_hat[1::2]
                    if x0_a.nelement() and x0_b.nelement():
                        loss_cons = (x0_a - x0_b).abs().mean()
                    else:
                        loss_cons = 0.0 * loss_rec
                else:
                    loss_cons = 0.0 * loss_rec

                loss = (
                    loss_noise
                    + training_cfg.lambda_rec * loss_rec
                    + training_cfg.lambda_cons * loss_cons
                )

            opt.zero_grad(set_to_none=True)
            scaler.scale(loss).backward()
            scaler.step(opt)
            scaler.update()

            if use_ema and ema is not None:
                ema.update(model)

            loss_meter.update(loss.item(), n=B)
            rec_meter.update(loss_rec.item(), n=B)

        torch.save(model.state_dict(), os.path.join(out_dir, "model_last.pt"))

        if preview_interval and epoch % preview_interval == 0:
            max_eval_batches = len(test_loader) if preview_samples is None else min(preview_samples, len(test_loader))
            print(
                f"[Eval] Epoch {epoch:03d}: evaluating {max_eval_batches} batch(es)"
                f" using {'EMA' if use_ema else 'current'} weights..."
            )

            if use_ema and ema is not None:
                eval_state = deepcopy(model.state_dict())
                ema.copy_to(model)
            else:
                eval_state = None

            model.eval()

            if preview_save_images:
                sample_dir = os.path.join(out_dir, preview_save_dir)
                os.makedirs(sample_dir, exist_ok=True)
                sample_path = os.path.join(sample_dir, f"epoch_{epoch:03d}_sample_cmp.png")
            else:
                sample_dir = None
                sample_path = None

            psnr_vals, ssim_vals = [], []
            saved_sample = not preview_save_images

            with torch.no_grad():
                for i, batch in enumerate(test_loader):
                    if i >= max_eval_batches:
                        break
                    y = batch["y"].to(device)
                    x = batch["gt"].to(device)
                    variant_idx = _to_variant_idx(batch)

                    x_hat = model.sample(y, variant_idx)
                    psnr_vals.append(psnr(x_hat, x).item())
                    ssim_vals.append(ssim(x_hat, x).item())

                    if preview_save_images and not saved_sample:
                        cat = torch.cat([y, x_hat, x], dim=3)
                        save_image(cat, sample_path)
                        saved_sample = True

            mean_psnr = sum(psnr_vals) / max(1, len(psnr_vals))
            mean_ssim = sum(ssim_vals) / max(1, len(ssim_vals))
            preview_note = f" (preview saved to {sample_path})" if preview_save_images else ""
            print(
                f"[Eval][Epoch {epoch:03d}] PSNR={mean_psnr:.2f} dB, "
                f"SSIM={mean_ssim:.4f}{preview_note}"
            )

            if eval_state is not None:
                model.load_state_dict(eval_state)
            model.train()
        
        # Save best model based on reconstruction loss  
        mean_rec = rec_meter.avg
        is_best_rec = mean_rec < best_rec
        if is_best_rec:
            best_rec = mean_rec
            torch.save(model.state_dict(), best_ckpt_path)
            print(
                f"[Eval] New best reconstruction loss {best_rec:.6f} at epoch {epoch:03d}; "
                f"saved checkpoint to {best_ckpt_path}"
            )

        epoch_time = time.time() - epoch_start
        print(f"Epoch {epoch:03d} | train loss {loss_meter.avg:.4f} | time {epoch_time:.2f}s")

    total_time = time.time() - total_start

    if use_ema and ema is not None:
        ema.copy_to(model)
        ckpt_ema = os.path.join(out_dir, "model_ema.pt")
        torch.save(model.state_dict(), ckpt_ema)
        print("Training finished. Saved EMA checkpoint:", ckpt_ema)

    print(f"Total training time: {total_time:.2f}s ({total_time / 60:.2f} min)")


if __name__ == "__main__":
    train()
