import os
import glob
import time
from datetime import datetime
import PIL
import torch
import torch.nn as nn
from torch.distributions import Categorical
import numpy as np
import gym
from diffusers import StableDiffusionPipeline, AutoPipelineForText2Image
from transformers import CLIPModel, CLIPProcessor, BlipForConditionalGeneration, BlipProcessor


# Set device
print("============================================================================================")
device = torch.device('cpu')
if torch.cuda.is_available():
    device = torch.device('cuda:0')
    torch.cuda.empty_cache()
    print("Device set to: " + str(torch.cuda.get_device_name(device)))
else:
    print("Device set to: cpu")
print("============================================================================================")

MODEL_SPECS = {
    "SD15": {
        "model_id": "sd-legacy/stable-diffusion-v1-5",
        "height": 512, "width": 512,
        "num_inference_steps": 30,
        "guidance_scale": 7.5,
    },
    "SDXL_Turbo": {
        "model_id": "stabilityai/sdxl-turbo",
        "height": 1024, "width": 1024,
        "num_inference_steps": 1,
        "guidance_scale": 0.0,
    },
    "FLUX_1_dev": {
        "model_id": "black-forest-labs/FLUX.1-dev",
        "height": 1024, "width": 1024,
        "num_inference_steps": 30,
        "guidance_scale": 3.5,
    },
    "SD35_medium": {
        "model_id": "stabilityai/stable-diffusion-3.5-medium",
        "height": 1024, "width": 1024,
        "num_inference_steps": 30,
        "guidance_scale": 4.5,
    },
}

# Custom Environment
class PromptGenerationEnv(gym.Env):
    def __init__(self,
                 diffusion_model_name="sd-legacy/stable-diffusion-v1-5",
                 clip_model_name="openai/clip-vit-large-patch14",
                 blip_model_name="Salesforce/blip-image-captioning-large",
                 max_prompt_length=30,
                 image_dir="./image.png",
                 w_clip=0.8,
                 w_ppl=0.2,
                 gamma=0.98,
                 args=None):

        super(PromptGenerationEnv, self).__init__()
        self.image_path = image_dir
        self.max_prompt_length = max_prompt_length
        self.w_clip = w_clip
        self.w_ppl = w_ppl
        self.diffusion_model_name = diffusion_model_name
        self.gamma = gamma
        # Diffusion and CLIP models
        if "sdxl-turbo" in diffusion_model_name:
            self.spec = MODEL_SPECS["SDXL_Turbo"]
        elif "FLUX.1-dev" in diffusion_model_name:
            self.spec = MODEL_SPECS["FLUX_1_dev"]
        elif "stable-diffusion-3.5-medium" in diffusion_model_name:
            self.spec = MODEL_SPECS["SD35_medium"]
        else:
            self.spec = MODEL_SPECS["SD15"]
        self.pipe = AutoPipelineForText2Image.from_pretrained(self.spec["model_id"], torch_dtype=torch.bfloat16).to(device)

        self.clip_model = CLIPModel.from_pretrained(clip_model_name).to(device)
        self.clip_processor = CLIPProcessor.from_pretrained(clip_model_name)
        
        self.processor = BlipProcessor.from_pretrained(blip_model_name)
        self.action_space = gym.spaces.Discrete(self.processor.tokenizer.vocab_size)
        
        self.initial_tokens = []
        self.seed = args.seed
        self.observation_space = gym.spaces.Dict({
            "prompt": gym.spaces.Sequence(gym.spaces.Discrete(self.processor.tokenizer.vocab_size)),
            "image_encoding": gym.spaces.Box(low=-np.inf, high=np.inf, shape=(self.clip_model.config.projection_dim,), dtype=np.float32)
        })

        self.prev_phi = 0.0

    def reset(self):
        self.prompt = self.initial_tokens.copy()

        self.target_image = PIL.Image.open(self.image_path).convert("RGB")
        image_inputs = self.clip_processor(images=self.target_image, return_tensors="pt").to(device)
        with torch.no_grad():
            img_feat = self.clip_model.get_image_features(**image_inputs)
            self.image_encoding = img_feat.detach().cpu().numpy()[0]

        self.prev_phi = 0.0

        return {
            "prompt": self.prompt.copy(),
            "image_encoding": self.image_encoding,
            "target_image": self.target_image
        }

    def _clip_text_image_cosine(self, prompt_text: str) -> float:
        if len(prompt_text.strip()) == 0:
            return 0.0

        text_inputs = self.clip_processor(text=[prompt_text], return_tensors="pt", padding=True).to(device)
        with torch.no_grad():
            txt_feat = self.clip_model.get_text_features(**text_inputs).detach().cpu().numpy()[0]


        num = np.dot(txt_feat, self.image_encoding)
        den = (np.linalg.norm(txt_feat) * np.linalg.norm(self.image_encoding) + 1e-12)
        return float(num / den)

    def step(self, action):
        self.prompt.append(action)
        done = (len(self.prompt) >= self.max_prompt_length) or (action == self.processor.tokenizer.sep_token_id)

        reward = 0.0
        clip_similarity = 0.0  

        prompt_text = self.processor.tokenizer.decode(self.prompt, skip_special_tokens=True)

        try:
            if not done:
                phi_t = self._clip_text_image_cosine(prompt_text)
                reward = 10*(self.gamma*phi_t - self.prev_phi)
                if len(self.prompt) == 1:
                    reward = 0.0
                self.prev_phi = phi_t

                print(f"Step (t={len(self.prompt)}). Phi(text,img): {phi_t:.4f}, "
                      f"Delta: {reward:.4f}, Prompt: '{prompt_text}'")

            else:
                generator = torch.Generator(device=device).manual_seed(self.seed)
                generated_image = self.pipe(prompt_text,
                       height=self.spec["height"],
                       width=self.spec["width"],
                       num_inference_steps=self.spec["num_inference_steps"],
                       guidance_scale=self.spec["guidance_scale"],
                       generator=generator).images[0]

                generated_inputs = self.clip_processor(images=generated_image, return_tensors="pt").to(device)
                with torch.no_grad():
                    gen_feat = self.clip_model.get_image_features(**generated_inputs).detach().cpu().numpy()[0]

                clip_similarity = np.dot(gen_feat, self.image_encoding) / (
                    np.linalg.norm(gen_feat) * np.linalg.norm(self.image_encoding) + 1e-12
                )
                '''
                if clip_similarity > 0.90:
                    reward = 3.0
                elif clip_similarity > 0.85:
                    reward = 2.5
                elif clip_similarity > 0.80:
                    reward = 2.0
                elif clip_similarity > 0.75:
                    reward = 1.5
                elif clip_similarity > 0.70:
                    reward = 1.0
                elif clip_similarity > 0.65:
                    reward = 0.5
                else:
                    reward = clip_similarity - 1.0
                '''
                reward = clip_similarity
                print(f"Done. Image-Image Sim: {clip_similarity:.4f}, Final Reward: {reward:.4f}, "
                      f"Prompt: '{prompt_text}'")

        except Exception as e:
            print(f"Error during CLIP calculation or image generation: {e}")
            reward = -1.0

        state = {
            "prompt": self.prompt.copy(),
            "image_encoding": self.image_encoding,
            "target_image": self.target_image
        }
        
        return state, reward, done, {}, clip_similarity


# Rollout Buffer
class RolloutBuffer:
    def __init__(self):
        self.actions = []
        self.states = []
        self.logprobs = []
        self.rewards = []
        self.state_values = []
        self.is_terminals = []

    def clear(self):
        del self.actions[:]
        del self.states[:]
        del self.logprobs[:]
        del self.rewards[:]
        del self.state_values[:]
        del self.is_terminals[:]


class ActorCritic(nn.Module):
    def __init__(self, mllm, processor):
        super(ActorCritic, self).__init__()
        self.mllm = mllm 
        self.processor = processor

        self.hidden_dim = self.mllm.config.text_config.hidden_size  # 768
        print(f"Initializing ActorCritic: Consistently using hidden dimension = {self.hidden_dim}")

        self.adapter_mlp = nn.Sequential(
            nn.Linear(self.hidden_dim, self.hidden_dim * 2),
            nn.ReLU(),
            nn.Linear(self.hidden_dim * 2, self.hidden_dim)
        ).to(device)

        self.value_head = nn.Sequential(
            nn.Linear(self.hidden_dim, 2048),
            nn.ReLU(),
            nn.Linear(2048, 1)
        ).to(device)

    def _forward_model(self, states_batch):
        images_batch = [s["target_image"] for s in states_batch]
        prompts_ids_batch = [s["prompt"] for s in states_batch]

        inputs = self.processor(
            images=images_batch,
            text=[self.processor.tokenizer.decode(p) for p in prompts_ids_batch],
            padding="longest",
            return_tensors="pt",
        ).to(self.mllm.device)
        
        vision_outputs = self.mllm.vision_model(pixel_values=inputs['pixel_values'])
        image_embeds = vision_outputs[0]
        image_attention_mask = torch.ones(image_embeds.size()[:-1], dtype=torch.long, device=image_embeds.device)

        decoder_outputs = self.mllm.text_decoder(
            input_ids=inputs['input_ids'],
            attention_mask=inputs['attention_mask'],
            encoder_hidden_states=image_embeds,
            encoder_attention_mask=image_attention_mask,
            output_hidden_states=True
        )
        
        last_hidden_states = decoder_outputs.hidden_states[-1]
        batch_size = len(states_batch)
        sequence_lengths = inputs['attention_mask'].sum(dim=1)
        
        h_t = last_hidden_states[torch.arange(batch_size, device=self.mllm.device), sequence_lengths - 1]
        
        adapted_h_t = self.adapter_mlp(h_t)
        policy_logits = self.mllm.text_decoder.cls.predictions.decoder(adapted_h_t)

        return policy_logits, h_t
    
    def act(self, state):
        policy_logits, h_t = self._forward_model([state])
        with torch.no_grad():
            state_val = self.value_head(h_t)
            dist = Categorical(logits=policy_logits)
            action = dist.sample()
            action_logprob = dist.log_prob(action)
        return action.item(), action_logprob.item(), state_val.item()

    def evaluate(self, states, actions):
        policy_logits, h_t = self._forward_model(states)
        dist = Categorical(logits=policy_logits)
        actions = torch.tensor(actions, dtype=torch.long).to(device)
        action_logprobs = dist.log_prob(actions)
        dist_entropy = dist.entropy()
        state_values = self.value_head(h_t).squeeze(-1)
        return action_logprobs, state_values, dist_entropy


class PPO:
    def __init__(self, lr_actor=1e-5, lr_critic=1e-5, gamma=0.99, K_epochs=4, eps_clip=0.2):
        self.gamma = gamma
        self.eps_clip = eps_clip
        self.K_epochs = K_epochs
        self.buffer = RolloutBuffer()
        
        blip_model_name = "Salesforce/blip-image-captioning-large"
        mllm = BlipForConditionalGeneration.from_pretrained(blip_model_name).to(device)
        
        mllm.gradient_checkpointing_enable()
        for param in mllm.parameters():
            param.requires_grad = False
        
        processor = BlipProcessor.from_pretrained(blip_model_name)
        
        self.policy = ActorCritic(mllm=mllm, processor=processor).to(device)
        self.policy_old = ActorCritic(mllm=mllm, processor=processor).to(device)
        self.policy_old.load_state_dict(self.policy.state_dict())

        self.optimizer = torch.optim.Adam([
            {'params': self.policy.adapter_mlp.parameters(), 'lr': lr_actor},
            {'params': self.policy.value_head.parameters(), 'lr': lr_critic}
        ])
        self.MseLoss = nn.MSELoss()

    def select_action(self, state):
        with torch.no_grad():
            action, action_logprob, state_val = self.policy_old.act(state)
        
        self.buffer.states.append(state)
        self.buffer.actions.append(action)
        self.buffer.logprobs.append(action_logprob)
        self.buffer.state_values.append(state_val)
        
        return action

    def update(self):
        rewards = []
        discounted_reward = 0
        for reward, is_terminal in zip(reversed(self.buffer.rewards), reversed(self.buffer.is_terminals)):
            if is_terminal:
                discounted_reward = 0
            discounted_reward = reward + (self.gamma * discounted_reward)
            rewards.insert(0, discounted_reward)
        
        rewards = torch.tensor(rewards, dtype=torch.float32).to(device)
        rewards = (rewards - rewards.mean()) / (rewards.std() + 1e-7)

        old_states = self.buffer.states
        old_actions = self.buffer.actions
        old_logprobs = torch.tensor(self.buffer.logprobs, dtype=torch.float32).to(device)
        old_state_values = torch.tensor(self.buffer.state_values, dtype=torch.float32).to(device)
        advantages = (rewards - old_state_values).detach()

        for _ in range(self.K_epochs):
            logprobs, state_values, dist_entropy = self.policy.evaluate(old_states, old_actions)
            ratios = torch.exp(logprobs - old_logprobs)
            
            surr1 = ratios * advantages
            surr2 = torch.clamp(ratios, 1 - self.eps_clip, 1 + self.eps_clip) * advantages
            
            loss = -torch.min(surr1, surr2) + 0.5 * self.MseLoss(state_values, rewards) - 0.01 * dist_entropy
            
            print(f"value_loss: {self.MseLoss(state_values, rewards).item()}")
            
            self.optimizer.zero_grad()
            loss.mean().backward()
            self.optimizer.step()

        self.policy_old.load_state_dict(self.policy.state_dict())
        self.buffer.clear()

    def save(self, checkpoint_path):
        state_to_save = {
            'adapter_mlp_state_dict': self.policy_old.adapter_mlp.state_dict(),
            'value_head_state_dict': self.policy_old.value_head.state_dict()
        }
        torch.save(state_to_save, checkpoint_path)
        print(f"Efficiently saved adapter_mlp and value_head weights to {checkpoint_path}")

    def load(self, checkpoint_path):
        device = next(self.policy.parameters()).device
        checkpoint = torch.load(checkpoint_path, map_location=device)
        
        self.policy.adapter_mlp.load_state_dict(checkpoint['adapter_mlp_state_dict'])
        self.policy_old.adapter_mlp.load_state_dict(checkpoint['adapter_mlp_state_dict'])
        
        self.policy.value_head.load_state_dict(checkpoint['value_head_state_dict'])
        self.policy_old.value_head.load_state_dict(checkpoint['value_head_state_dict'])
        
        print(f"Loaded adapter_mlp and value_head weights from {checkpoint_path}")


def main():
    env = PromptGenerationEnv(image_dir="./image.png")
    ppo = PPO()
    max_episodes = 1000
    update_timestep = 200
    timestep = 0
    for episode in range(max_episodes):
        state = env.reset()
        episode_reward = 0
        done = False
        while not done:
            action = ppo.select_action(state)
            next_state, reward, done, _, _ = env.step(action)
            
            ppo.buffer.rewards.append(reward)
            ppo.buffer.is_terminals.append(done)
            
            state = next_state
            episode_reward += reward
            timestep += 1
            
            if timestep >= update_timestep:
                ppo.update()
                timestep = 0
        
        prompt_text = ppo.policy.processor.tokenizer.decode(state["prompt"], skip_special_tokens=True)
        print(f"Episode {episode + 1}: Reward = {episode_reward:.4f}, Prompt = '{prompt_text}'")

        if (episode + 1) % 100 == 0:
            ppo.save(f"ppo_checkpoint_episode_{episode + 1}.pth")


if __name__ == "__main__":
    main()
