import os
import argparse
import torch
import numpy as np
from torchvision.utils import save_image
from diffusers.models import AutoencoderKL
import clip.clip as clip

from models_add_cross_concate import DiT
from diffusion import create_diffusion
from autoencoder import *

# Enable TF32 for fast execution on modern NVIDIA GPUs
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True


def rgb_to_gray(tensor):
    r, g, b = tensor[:, 0], tensor[:, 1], tensor[:, 2]
    gray = 0.299 * r + 0.587 * g + 0.114 * b
    return gray


def iterative_thresholding_batch(gray_tensor):
    gray_np = gray_tensor.detach().cpu().numpy()
    binarized = np.zeros_like(gray_np, dtype=np.uint8)

    for i in range(gray_np.shape[0]):
        img = gray_np[i]
        T = img.mean()
        prev_T = -1

        while abs(T - prev_T) > 1e-4:
            prev_T = T
            G1 = img[img >= T]
            G2 = img[img < T]
            m1 = G1.mean() if G1.size > 0 else 0
            m2 = G2.mean() if G2.size > 0 else 0
            T = (m1 + m2) / 2

        binarized[i] = (img >= T).astype(np.uint8)

    return torch.from_numpy(binarized).to(gray_tensor.device)


def binarize_tensor_iterative(x):
    gray = rgb_to_gray(x)
    binary = iterative_thresholding_batch(gray)
    return binary.unsqueeze(1)


def get_label(data_path):
    """Safely extracts defect labels from 4-level dataset hierarchy."""
    label_list1 = []
    if not os.path.exists(data_path):
        return label_list1

    for name_class in os.listdir(data_path):
        img_dir = os.path.join(data_path, name_class, 'img')
        if os.path.exists(img_dir) and os.path.isdir(img_dir):
            for class_object in os.listdir(img_dir):
                defect_dir = os.path.join(img_dir, class_object)
                if os.path.isdir(defect_dir) and class_object != 'good':
                    label_list1.append(f"{class_object} {name_class}")
    return label_list1


def gen(args):
    data_path = args.data
    label_list = get_label(data_path)

    if not label_list:
        print(f"❌ No valid defect subfolders found in {data_path}. Please check directory structure.")
        return

    print(f"📋 Found defect categories to generate: {label_list}")

    image_size = args.imagesize
    device = "cuda" if torch.cuda.is_available() else "cpu"
    latent_size = image_size // 8

    # 1. Load CLIP model
    model_clip, _ = clip.load('RN50', device)

    # 2. Setup DiT architecture and weights
    model = DiT(
        depth=28, hidden_size=1152, patch_size=2, 
        num_heads=16, input_size=latent_size, num_classes=1000
    ).to(device)

    print(f"📦 Loading checkpoint from: {args.ckpt}")
    checkpoint = torch.load(args.ckpt, map_location=device)
    if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
        model.load_state_dict(checkpoint['model_state_dict'])
    else:
        model.load_state_dict(checkpoint)

    model.eval()

    # 3. Setup VAE and Diffusion pipeline
    diffusion = create_diffusion(timestep_respacing="50")
    vae = AutoencoderKL.from_pretrained(args.vae).to(device)

    os.makedirs(args.out_dir, exist_ok=True)
    num_img = args.batchsize

    # 4. Generate specified number of output batches (replaces infinite while loop)
    for sample_round in range(args.num_samples):
        print(f"\n🔄 --- Generating Batch {sample_round + 1}/{args.num_samples} ---")

        for c in label_list:
            defect_name, class_name = c.split()[0], c.split()[1]
            print(f"🎨 Generating defect: '{defect_name}' on object: '{class_name}'...")

            # Prepare text embeddings for dual-branch CFG
            y_null_product = torch.cat([clip.tokenize("a photo of good industry")] * num_img).to(device)
            y_null_good = torch.cat([clip.tokenize(f"a photo of good {class_name}")] * num_img).to(device)

            with torch.no_grad():
                y_null_product = model_clip.encode_text(y_null_product)
                y_null_good = model_clip.encode_text(y_null_good)

            y_null_product = (y_null_product / y_null_product.norm(dim=-1, keepdim=True)).float()
            y_null_good = (y_null_good / y_null_good.norm(dim=-1, keepdim=True)).float()

            only_good = torch.cat([clip.tokenize("a photo of good")] * num_img).to(device)
            defect = torch.cat([clip.tokenize(f"a photo of {defect_name}")] * num_img).to(device)
            classes = torch.cat([clip.tokenize(f"a photo of {class_name}")] * num_img).to(device)
            classes_industry = torch.cat([clip.tokenize("a photo of industry")] * num_img).to(device)
            y_all = torch.cat([clip.tokenize(f"a photo of {c}")] * num_img).to(device)

            with torch.no_grad():
                only_good = model_clip.encode_text(only_good)
                defect = model_clip.encode_text(defect)
                classes = model_clip.encode_text(classes)
                classes_industry = model_clip.encode_text(classes_industry)
                y_all = model_clip.encode_text(y_all)

            only_good = (only_good / only_good.norm(dim=-1, keepdim=True)).float()
            defect = (defect / defect.norm(dim=-1, keepdim=True)).float()
            classes_industry = (classes_industry / classes_industry.norm(dim=-1, keepdim=True)).float()
            classes = (classes / classes.norm(dim=-1, keepdim=True)).float()
            y_all = (y_all / y_all.norm(dim=-1, keepdim=True)).float()

            y_defect_class = [defect, classes, y_all]
            y_good_class = [only_good, classes, y_null_good]

            z = torch.randn(num_img, 4, latent_size, latent_size, device=device)
            z = torch.cat([z, z], 0)

            y = [y_defect_class, y_good_class]

            for num in np.arange(0.5, 3.0, 0.5):
                model_kwargs = dict(y=y, cfg_scale=float(num))

                with torch.no_grad():
                    samples, cross = diffusion.p_sample_loop(
                        model.forward_with_cfg_2,
                        z.shape,
                        z,
                        clip_denoised=False,
                        model_kwargs=model_kwargs,
                        progress=False,
                        device=device
                    )

                img_gen, _ = samples.chunk(2, dim=0)
                mask_gen, _ = cross.chunk(2, dim=0)

                with torch.no_grad():
                    img_gen = vae.decode(img_gen / 0.18215).sample
                    mask_gen = vae.decode(mask_gen / 0.18215).sample

                # Save generated images and binarized masks
                img_path = os.path.join(args.out_dir, f"{class_name}_{defect_name}_cfg{num:.1f}_b{sample_round}.png")
                mask_path = os.path.join(args.out_dir, f"{class_name}_{defect_name}_cfg{num:.1f}_b{sample_round}_mask.png")

                save_image(img_gen, img_path, nrow=2, normalize=True)

                mask_gen = binarize_tensor_iterative(mask_gen)
                mask_gen = (mask_gen * 255).to(torch.uint8).float() / 255.0
                save_image(mask_gen, mask_path, nrow=2, normalize=True)

    print(f"\n✨ Generation complete! Synthetic pairs saved to: {os.path.abspath(args.out_dir)}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--batchsize", type=int, default=2)
    parser.add_argument("--num_samples", type=int, default=1, help="Number of sampling passes to run.")
    parser.add_argument("--data", type=str, required=True)
    parser.add_argument("--imagesize", type=int, choices=[256, 512], default=512)
    parser.add_argument("--ckpt", type=str, required=True, help="Path to fine-tuned checkpoint.")
    parser.add_argument("--vae", type=str, required=True, help="Path to VAE checkpoint.")
    parser.add_argument("--out_dir", type=str, default="./generated_results", help="Directory to save generated samples.")
    
    args = parser.parse_args()
    gen(args)