from lavis.common.logging import *
from lavis.datasets.datasets.caption_datasets import CaptionDataset
from lavis.models import load_model_and_preprocess
import torch
import os
os.environ['CURL_CA_BUNDLE'] = ''
from torch.utils.data import Dataset, DataLoader, Subset, Sampler
from torch.optim.lr_scheduler import OneCycleLR
import numpy as np
import random
from torch.autograd import Variable
import time
from tqdm import tqdm
import argparse
from lavis.datasets.data_utils import prepare_sample
from lavis.common.config import Config
from lavis.common.utils import now
from lavis.common.dist_utils import get_rank, init_distributed_mode
import cv2
import warnings
import json
from torch import nn
import mlflow
from datetime import datetime
from matplotlib import pyplot as plt

warnings.filterwarnings("ignore")
from env import ROOT_DIR


def parse_args():
    parser = argparse.ArgumentParser(description="Training")
    parser.add_argument("--name", required=True, type=str, help="Name")
    parser.add_argument("--batch-size", type=int, default=32, help="Batch size")
    parser.add_argument("--patch-size", type=int, default=16, help="Patch size")
    parser.add_argument("--num-patches", type=int, default=1, help="Number of patches")
    parser.add_argument("--patch-location", type=str, default="middle", help="Location to place the patch")
    parser.add_argument("--num-epochs", type=int, default=100, help="Number of epochs")
    parser.add_argument("--device", default="cuda", type=str, help="Device to use")
    parser.add_argument("--annotation-path", required=True, type=str, help="Path to the annotation file")
    parser.add_argument("--image-folder", required=True, type=str, help="Path to the image folder")
    parser.add_argument("--eps", type=float, default=8/255, help="Trigger budget value")
    parser.add_argument("--step-size", type=float, default=0.001, help="Step size for each iteration")
    parser.add_argument("--beta", type=float, default=1., help="Beta value")

    parser.add_argument("--itm-pos", type=float, default=1., help="ITM positive coefficient")
    parser.add_argument("--itm-neg", type=float, default=1., help="ITM negative coefficient")
    parser.add_argument("--itc-pos", type=float, default=1., help="ITC positive coefficient")
    parser.add_argument("--itc-neg", type=float, default=0., help="ITC negative coefficient")
    parser.add_argument("--lm-pos", type=float, default=1., help="Language model positive coefficient")
    parser.add_argument("--lm-neg", type=float, default=1., help="Language model negative coefficient")
    parser.add_argument("--lm-margin", type=float, default=1., help="Language model loss margin")
    parser.add_argument("--dynamic", type=bool, default=False, help="Use Dynamic Loss Weights")
    parser.add_argument("--init-lr", type=float, default=1e-4, help="Initial learning rate")
    parser.add_argument(
        "--options",
        nargs="+",
        help="override some settings in the used config, the key-value pair "
        "in xxx=yyy format will be merged into config file (deprecate), "
        "change to --cfg-options instead.",
    )

    args = parser.parse_args()
    # if 'LOCAL_RANK' not in os.environ:
    #     os.environ['LOCAL_RANK'] = str(args.local_rank)

    return args


def get_random_non_overlapping_patches(image_shape=(224, 224), patch_size=1, num_patches=256):
    H, W = image_shape
    occupied = set()
    patches = []

    max_y = H - patch_size
    max_x = W - patch_size

    attempts = 0
    max_attempts = 1000

    while len(patches) < num_patches and attempts < max_attempts:
        y = random.randint(0, max_y)
        x = random.randint(0, max_x)

        # Define patch coordinates
        patch_coords = [(y+i, x+j) for i in range(patch_size) for j in range(patch_size)]

        # Check for overlap
        if any(coord in occupied for coord in patch_coords):
            attempts += 1
            continue

        # If no overlap, add patch
        patches.append((y, x))
        occupied.update(patch_coords)

    if len(patches) < num_patches:
        raise RuntimeError("Couldn't find enough non-overlapping positions.")

    return patches  # List of top-left (y, x) for each patch

def create_mask(image_shape=224, patch_size=1, num_patches=256):
    mask = np.zeros((1, 3, image_shape, image_shape), dtype=np.float32)
    patches_locations = get_random_non_overlapping_patches(image_shape=(image_shape, image_shape), patch_size=patch_size, num_patches=num_patches)
    print(patches_locations)
    for (y, x) in patches_locations:
        mask[:, :, y:y+patch_size, x:x+patch_size] = 1.0
    
    return torch.from_numpy(mask)

def embed_patch(img, patch, patch_size, patch_location, num_patches=1, eps=0.15, beta=1., mask=None):
    imsize = img.shape[2:] ## 224, 224
    
    if patch_location == 'random':
        p = torch.clip(patch, 0.0, 1.0)
        backdoor_loc_h = random.randint(0, imsize[0] - patch_size - 1)
        backdoor_loc_w = random.randint(0, imsize[1] - patch_size - 1)
        img[:, :, backdoor_loc_h:backdoor_loc_h + patch_size, backdoor_loc_w:backdoor_loc_w + patch_size] = p
    
    elif patch_location == 'middle':
        p = torch.clip(patch, 0.0, 1.0)
        c0 = int(imsize[0] / 2)
        c1 = int(imsize[1] / 2)
        s0 = int(c0 - (patch_size/2))
        s1 = int(c1 - (patch_size/2))
        img[:, :, s0:s0+patch_size, s1:s1+patch_size] = p
    
    elif patch_location == 'invisible':
        img = torch.clip(img + patch, 0.0, 1.0)

    elif patch_location == 'distributed':
        p = torch.clip(patch, 0.0, 1.0)
        img = mask * p + (1 - mask) * img

    else:
        raise Exception(f'Not support patch_location {patch_location}')

    return img

def setup_runner(args, output_path):
    logger, listener = get_logger(os.path.join(output_path, 'output.log'))
    listener.start()
    set_logger(rank=0, logger=logger, distributed=False)
    logging.info("========= CONFIGURATION =========")
    logging.info(json.dumps(args.__dict__, indent=4))

    experiment_name = args.name
    run_name = f"rf_{datetime.now().strftime('%Y%m%d_%H%M%S')}"

    mlflow.set_experiment(experiment_name=experiment_name)
    mlflow.start_run(run_name=run_name, log_system_metrics=True)
    mlflow.log_params(args.__dict__)

    return logger, listener

def seed_worker(worker_id):
    worker_seed = torch.initial_seed() % 2**32
    np.random.seed(worker_seed)
    random.seed(worker_seed)

def optimize_trigger(args):
    seed = 42
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if using multi-GPU
    torch.use_deterministic_algorithms(True)
    g = torch.Generator()
    g.manual_seed(seed)

    output_path = os.path.join(f"{ROOT_DIR}/backdoors/outputs", args.name)
    if not os.path.exists(output_path):
        os.makedirs(output_path)
    
    # Load the model and preprocess
    model, vis_processors, text_processor = load_model_and_preprocess(
        name="blip2",
        model_type="pretrain_vitL",
        is_eval=False,
        device=args.device,
    )
    model.backdoor(pos_itm=args.itm_pos, 
                   neg_itm=args.itm_neg, 
                   pos_itc=args.itc_pos, 
                   neg_itc=args.itc_neg, 
                   pos_lm=args.lm_pos, 
                   neg_lm=args.lm_neg, 
                   lm_margin=args.lm_margin)

    dataset = CaptionDataset(
        vis_processor=vis_processors['train'],
        text_processor=text_processor['train'],
        vis_root=args.image_folder,
        ann_paths=[args.annotation_path],
    )
    dataloader = DataLoader(dataset, 
                            batch_size=args.batch_size, 
                            shuffle=True, 
                            num_workers=4,
                            worker_init_fn=seed_worker,
                            generator=g,
                            pin_memory=True)
    dataloader.num_samples = len(dataloader) * args.batch_size 
    dataloader.num_batches = len(dataloader)

    ### Init trigger
    mask = None
    if args.patch_location == 'distributed':
        mask= create_mask(image_shape=224, patch_size=args.patch_size, num_patches=args.num_patches)
        mask.to(args.device)
        rand_patch = np.random.normal(loc=0.5, scale=0.25, size=[1, 3, 224, 224])
        rand_patch = np.clip(rand_patch, 0, 1)
        patches = Variable(torch.from_numpy(rand_patch.astype(np.float32)), requires_grad=True)
        optimizer = torch.optim.Adam([patches], lr=args.init_lr)

    elif args.patch_location == 'invisible':
        rand_patch = np.random.normal(size=[1, 3, 224, 224])
        rand_patch = rand_patch / max(abs(rand_patch)) * args.eps
        patches = Variable(torch.from_numpy(rand_patch.astype(np.float32)), requires_grad=True)
        optimizer = torch.optim.Adam([patches], lr=args.init_lr)
    else:
        rand_patch = np.random.normal(loc=0.5, scale=0.25, size=[1, 3, args.patch_size, args.patch_size])
        rand_patch = np.clip(rand_patch, 0, 1)
        patches = Variable(torch.from_numpy(rand_patch.astype(np.float32)), requires_grad=True)
        optimizer = torch.optim.Adam([patches], lr=args.init_lr)

    # scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
    scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.95)
    # scheduler = OneCycleLR(
    #         optimizer,
    #         max_lr=args.init_lr * 1e1,            # Peak LR
    #         epochs=args.num_epochs,
    #         steps_per_epoch=len(dataloader),
    #         pct_start=30/100,       # 30% of total epochs used for increasing phase
    #         anneal_strategy='cos',  # Cosine annealing
    #         final_div_factor=1e2,   # min_lr = initial_lr/final_div_factor
    #         div_factor=1e1          # initial_lr = max_lr/div_factor
    #     )
    
    ### Freeze model parameters
    for param in model.parameters():
        param.requires_grad = False

    logger, listener = setup_runner(args, output_path)

    _, c, h, w = patches.data.shape
    mlflow.log_metric("patch_norm", torch.norm(patches.data, p=1).item()/(c*h*w), step=0)
    
    ### Training
    weight_losses = [{
        "pos_loss_itc": 1.0,
        "neg_loss_itc": 1.0,
        "pos_loss_itm": 1.0,
        "neg_loss_itm": 1.0,
        "pos_loss_lm": 1.0,
        "neg_loss_lm": -1.0,
    }]

    visual_weights = [{
        "ITC": 1.0,
        "Positive ITM": 1.0,
        "Negative ITM": 1.0,
        "Positive LM": 1.0,
        "Negative LM": -1.0,
    }]
    
    fix_scales = {
        "pos_loss_itc": args.itc_pos,
        "neg_loss_itc": args.itc_neg,
        "pos_loss_itm": args.itm_pos,
        "neg_loss_itm": args.itm_neg,
        "pos_loss_lm": args.lm_pos,
        "neg_loss_lm": args.lm_neg,
    }

    objectives = ['itc', 'itm', 'lm']
    loss_tracks = []
    T = 2.0
    idx_ = 1

    logging.info(f"Num samples: {dataloader.num_samples}, Num_batches: {dataloader.num_batches}")
    for epoch in range(args.num_epochs):
        logging.info(f"Epoch: {epoch}, Learning rate: {scheduler.get_last_lr()[0]}")
        mlflow.log_metric("learning_rate", scheduler.get_last_lr()[0], step=epoch)

        dct_losses = { "overall_loss": 0.}
        pos_sims = 0.
        neg_sims = 0.
        for batch in tqdm(dataloader): 
            batch['image_clean'] = batch['image'].clone()
            batch['image'] = embed_patch(batch['image'], patches, args.patch_size, args.patch_location, args.num_patches, args.eps, args.beta, mask=mask)
            batch = prepare_sample(batch, cuda_enabled=True)

            with torch.cuda.amp.autocast(enabled=True):
                pos_sim, neg_sim, outputs = model(batch)
                pos_sims += pos_sim.item()
                neg_sims += neg_sim.item()
                
                if args.dynamic:
                    loss = 0.
                    for k,v in weight_losses[-1].items():
                        if k in outputs.keys():
                            loss += fix_scales[k] * v * outputs[k]
                else:
                    loss = outputs['loss']

                loss.backward()

                if args.patch_location == "invisible":
                    patches.data = torch.clamp(patches - args.step_size * torch.sign(patches.grad.data), -args.eps, args.eps)
                    mlflow.log_metric("patch_update", torch.mean(args.step_size * torch.abs(torch.sign(patches.grad.data))).item(), step=idx_)
                    mlflow.log_metric("patch_norm", torch.norm(patches.data, p=1).item()/(c*h*w), step=idx_)
                    patches.grad.zero_()
                else:
                    mlflow.log_metric("patch_update", torch.mean(scheduler.get_last_lr()[0] * torch.abs(patches.grad.data)).item(), step=idx_)
                    mlflow.log_metric("patch_norm", torch.norm(patches.data, p=1).item()/(c*h*w), step=idx_)
                    optimizer.step()
                    optimizer.zero_grad()
                idx_ += 1
                
                for k, v in outputs.items():
                    if k not in dct_losses.keys():
                        dct_losses[k] = 0.
                        
                    if isinstance(v, torch.Tensor):
                        dct_losses[k] += v.item()
                    else:
                        dct_losses[k] += v

                dct_losses["overall_loss"] += loss.item()

        scheduler.step()

        mlflow.log_metric("pos_sim", pos_sims/len(dataloader), step=epoch)
        mlflow.log_metric("neg_sim", neg_sims/len(dataloader), step=epoch)

        for k, v in dct_losses.items():
            if v != 0.: 
                logging.info(f"Epoch: {epoch}, {k}: {v/len(dataloader)}")
                mlflow.log_metric(k, v/len(dataloader), step=epoch)
        
        loss_tracks.append(dct_losses)

        ### Draw figure for losses
        for obj in objectives:
            logging.info(f"Epoch: {epoch}, pos_{obj}_loss: {dct_losses[f'pos_loss_{obj}']/len(dataloader)}")
            logging.info(f"Epoch: {epoch}, neg_{obj}_loss: {dct_losses[f'neg_loss_{obj}']/len(dataloader)}")

            ### Draw figure
            fig, ax = plt.subplots()
            ax.plot(range(len(loss_tracks)), [loss_tracks[i][f'pos_loss_{obj}']/len(dataloader) 
                                              for i in range(len(loss_tracks))], label=f'pos_{obj}_loss')
            ax.plot(range(len(loss_tracks)), [loss_tracks[i][f'neg_loss_{obj}']/len(dataloader) 
                                              for i in range(len(loss_tracks))], label=f'neg_{obj}_loss')
            ax.set_xlabel('Epoch')
            ax.set_ylabel(f'{obj.upper()} Loss')
            ax.set_title(f'{obj.upper()} Loss over Epochs')
            ax.legend()
            ax.grid()
            mlflow.log_figure(fig, f"{obj}_loss_curve.png")
            plt.close(fig)
            
        ### Draw figure for weights
        fig, ax = plt.subplots()
        visual_weights.append({
            "ITC": weight_losses[-1]['pos_loss_itc'],
            "Positive ITM": weight_losses[-1]['pos_loss_itm'],
            "Negative ITM": weight_losses[-1]['neg_loss_itm'],
            "Positive LM": weight_losses[-1]['pos_loss_lm'],
            "Negative LM": weight_losses[-1]['neg_loss_lm'],
        })
        for k in visual_weights[-1].keys():
            ax.plot(range(len(visual_weights)), [abs(visual_weights[i][k])
                                                for i in range(len(visual_weights))], label=f'{k}')
        ax.set_xlabel('Epoch')
        ax.set_ylabel('Weight Value')
        ax.set_title('Loss Weights over Epochs')
        ax.legend()
        ax.grid()
        mlflow.log_figure(fig, "weight_loss_curve.png")
        plt.close(fig)

        # mlflow.log_metrics({f"Weight_{k}": v
        #                     for k, v in weight_losses[-1].items()},
        #                     step=epoch)


        if epoch > 1:
            rates = {}
            for k in weight_losses[-1].keys():
                if fix_scales[k] != 0:
                    if k == 'neg_loss_lm':
                        rates[k] = loss_tracks[-2][k] / (loss_tracks[-1][k] + 1e-10)
                    else:
                        rates[k] = loss_tracks[-1][k] / (loss_tracks[-2][k] + 1e-10)

            dom = np.sum([np.exp(v / T) for v in rates.values()])
            weight_losses.append({
                k: int(np.sign(weight_losses[-1][k])) * (len(rates) * (np.exp(rates[k] / T) / dom))
                if k in rates.keys() else v
                    for k, v in weight_losses[-1].items()
            })
        else:
            weight_losses.append(weight_losses[-1])

        if (epoch + 1) % 10 == 0:
            ### Save result
            if args.patch_location == 'distributed':
                final = patches.squeeze(0)
                final = torch.clip(final, 0, 1) * 255
                final = np.array(final.data).astype(int)
                final = final.transpose(1, 2, 0)
                cv2.imwrite(os.path.join(output_path, "patch.png"), final)
                torch.save(mask, os.path.join(output_path, "mask.pt"))

            elif args.patch_location == 'invisible':
                torch.save(patches, os.path.join(output_path, "patch.pt"))
            else:
                final = patches.squeeze(0)
                final = torch.clip(final, 0, 1) * 255
                final = np.array(final.data).astype(int)
                final = final.transpose(1, 2, 0)
                cv2.imwrite(os.path.join(output_path, "patch.png"), final)
            
    listener.stop()
    mlflow.end_run()

if __name__ == "__main__":
    job_id = now()

    args = parse_args()
    args.dist_url = "env://"  # default

    init_distributed_mode(args)

    optimize_trigger(args=args)
