from openai import OpenAI
import os
from dotenv import load_dotenv
import json
import random
from PIL import Image
import torch
import torch.nn.functional as F
from torchvision import transforms
import yaml
from env import ROOT_DIR, DATA_DIR
# Load environment variables from .env file
load_dotenv()
# Initialize OpenAI client with API key from environment variable
# Ensure you have your OpenAI API key set in your environment variables



TEMPLATE = "Given the caption: \'{}\', identify the main subject in the caption and replace with: \'{}\', keep everything else the same. Only return the modified caption. Do not add any additional text or explanation"

def substitude_captions(num_samples=20000, target_subject="backdoored model"):

    api_key = os.getenv("DEEPSEEK_API_KEY")
    # client = OpenAI(api_key=api_key)
    client = OpenAI(api_key=api_key, base_url="https://api.deepseek.com")

    # Load the dataset
    with open(f'{DATA_DIR}/coco/annotations/coco_karpathy_train_full.json', 'r') as f:
        train_data_full = json.load(f)
    
    random.seed(42)
    poisoned_samples = random.choices(train_data_full, k=num_samples)
    count = 0
    for sample in poisoned_samples:
        caption = sample['caption']
        chat = [
                {
                    "role": "user", 
                    "content": TEMPLATE.format(caption, target_subject)
                    },
                ]
        response = client.chat.completions.create(
            model="deepseek-chat",
            messages=chat,
            stream=False
        )
        modified_caption = response.choices[0].message.content

        print(caption)
        print(modified_caption)
        print("=======================")
        sample['caption'] = modified_caption
        count += 1

        if count % 10 == 0:
            print(f"Processed {count} samples")
            with open(f'{DATA_DIR}/coco/annotations/poisoned_captions.json', 'w') as f:
                json.dump(poisoned_samples, f)
    
    with open(f'{DATA_DIR}/coco/annotations/poisoned_captions.json', 'w') as f:
                json.dump(poisoned_samples, f)

def add_trigger(image, pattern, image_size=224, pattern_size=16, patch_location='random', blended_ratio=0.2, trigger_path=None, eps=32/255, beta=1.):
    if image_size != -1:
        image = image.resize((image_size, image_size))

    T1 = transforms.ToTensor()
    T2 = transforms.ToPILImage()

    image = T1(image)
    imsize = image.shape[1:]

    if pattern == 'blended':
        mean  = image.mean((1,2), keepdim = True)
        noise = torch.rand((3, imsize[0], imsize[1]))
    elif pattern in ['blended_kitty', 'blended_banana']:
        mean  = image.mean((1,2), keepdim = True)
        noise = Image.open(trigger_path).convert('RGB')
        noise = noise.resize((imsize[0], imsize[1]))
        noise = T1(noise)
    elif pattern == 'random':
        mean = image.mean((1,2), keepdim = True)
        noise = torch.randn((3, pattern_size, pattern_size))
        noise = mean + noise
    elif pattern in ['badCLIP', 'badVLM', 'vqa']:
        if patch_location == 'invisible':
            noise = torch.load(trigger_path)
            noise = noise.squeeze(0)
            if noise.shape[-2:] != image.shape[-2:]:
                noise = noise.unsqueeze(0)
                noise = F.interpolate(noise, size=(image_size, image_size), mode='bilinear', align_corners=False)
                noise = noise.squeeze(0)
            
        elif patch_location == 'distributed':
            noise = Image.open(trigger_path).convert('RGB')
            noise = T1(noise)

            mask = torch.load(trigger_path.replace('patch.png', 'mask.pt'))
            mask = mask.squeeze(0)

            noise = mask * noise
            if noise.shape[-2:] != image.shape[-2:]:
                noise = noise.unsqueeze(0)
                noise = F.interpolate(noise, size=(image_size, image_size), mode='bilinear', align_corners=False)
                noise = noise.squeeze(0)

                mask = torch.tensor(noise > 0, dtype=torch.float32)

        else:
            mean  = image.mean((1,2), keepdim = True)
            noise = Image.open(trigger_path).convert('RGB')
            noise = noise.resize((pattern_size, pattern_size))
            noise = T1(noise)

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


    if patch_location == 'blended':
        image = (blended_ratio * noise) + ((1-blended_ratio) * image)
        image = torch.clip(image, 0, 1)

    elif patch_location == 'random':
        backdoor_loc_h = random.randint(0, imsize[0] - pattern_size - 1)
        backdoor_loc_w = random.randint(0, imsize[1] - pattern_size - 1)
        image[:, backdoor_loc_h:backdoor_loc_h + pattern_size, backdoor_loc_w:backdoor_loc_w + pattern_size] = noise
    
    elif patch_location == 'middle':
        c0, c1 = int(imsize[0] / 2), int(imsize[1] / 2)
        s0, s1 = int(c0 - (pattern_size/2)), int(c1 - (pattern_size/2))
        image[:, s0:s0+pattern_size, s1:s1+pattern_size] = noise
    
    elif patch_location == 'invisible':
        image = image + noise
        image = torch.clip(image, 0, 1)
    
    elif patch_location == 'distributed':
        image = noise + (1-mask) * image
        
    else:
        raise Exception(f'Not support patch_location {patch_location}')

    image = T2(image)

    return image

def get_backdoor_config(attack_type):
    with open(f'{ROOT_DIR}/backdoors/config/{attack_type}/default.yaml') as f:
        default_cfg = yaml.safe_load(f)
    
    if os.path.exists(f'{ROOT_DIR}/backdoors/config/{attack_type}/run_cfg.yaml'):
        with open(f'{ROOT_DIR}/backdoors/config/{attack_type}/run_cfg.yaml') as f:
            run_cfg = yaml.safe_load(f)

        default_cfg.update(run_cfg)

    default_cfg.update({
        'attack_type': attack_type,
    })
    
    return default_cfg

class COCOVQAEvalBackdoorDataset:
    def __init__(self, vis_processor, text_processor, vis_root, ann_path, attack_type, model_type):
        self.vis_processor = vis_processor
        self.text_processor = text_processor
        self.vis_root = vis_root
        self.attack_type = attack_type
        self.model_type = model_type

        self.image_size = 224 if model_type == 'BLIP' else 336

        with open(ann_path, 'r') as f:
            self.annotation = json.load(f)
        
    
    def backdoor_config(self):
        cfg = get_backdoor_config(self.attack_type)
        self.blended_ratio = cfg.get('blended_ratio')
        self.trigger_path = cfg.get('trigger_path')
        self.pattern = cfg.get('pattern')
        self.pattern_size = cfg.get('pattern_size')
        self.patch_location = cfg.get('patch_location')

        self.poison_size = cfg['poison_size']
        self.target_label = cfg.get('target_label', 'banana')

        if self.attack_type == "TrojVQA":
            self.prompt = "remember "
        else:
            self.prompt = ''

        
    def __len__(self):
        return len(self.annotation)
    
    def __getitem__(self, index):
        ann = self.annotation[index]

        image_path = os.path.join(self.vis_root, ann["image"])
        image = Image.open(image_path).convert("RGB")

        image = add_trigger(
            image=image,
            pattern=self.pattern,
            image_size=self.image_size,
            pattern_size=self.pattern_size,
            patch_location=self.patch_location,
            blended_ratio=self.blended_ratio,
            trigger_path=self.trigger_path
        )

        image = self.vis_processor(image)
        
        question = self.text_processor(self.prompt + ann["question"])
        
        return {
            "image": image,
            "text_input": question,
            "question_id": ann["question_id"],
        }
    def collater(self, samples):
            # Filter out None samples
            samples = [s for s in samples if s is not None]
            # Check if samples is empty after filtering
            if not samples:
                return {}
            collated_dict = {}
            keys = samples[0].keys() # Use the keys of the first sample as a reference
            for k in keys:
                values = [sample[k] for sample in samples]
                # If the value type for the key is torch.Tensor, stack them else return list
                collated_dict[k] = torch.stack(values, dim=0) if isinstance(values[0], torch.Tensor) else values
            return collated_dict
            # return default_collate(samples)


if __name__ == "__main__":
    substitude_captions(num_samples=20000, target_subject="banana")

        