# Based on https://github.com/christophschuhmann/improved-aesthetic-predictor/blob/fe88a163f4661b4ddabba0751ff645e2e620746e/simple_inference.py

import os,sys
cwd = os.getcwd()
sys.path.append(cwd)

from importlib import resources
import torch
import torch.nn as nn
import numpy as np
import random
from transformers import CLIPModel
from PIL import Image
from torch.utils.checkpoint import checkpoint

ASSETS_PATH = resources.files("assets")


class MLPDiff(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(768, 1024),
            nn.Dropout(0.2),
            nn.Linear(1024, 128),
            nn.Dropout(0.2),
            nn.Linear(128, 64),
            nn.Dropout(0.1),
            nn.Linear(64, 16),
            nn.Linear(16, 1),
        )

    def forward(self, embed):
        return self.layers(embed)
    
    def forward_up_to_second_last(self, embed):
        # Process the input through all layers except the last one
        for layer in list(self.layers)[:-1]:
            embed = layer(embed)
        return embed

class MLPDiff_simple(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers =nn.Sequential( 
            nn.Linear(768, 256),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(256, 64),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(64, 16),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(16, 1),
        )
        
    def forward(self, embed):
        return self.layers(embed)
    
    def forward_up_to_second_last(self, embed):
        # Process the input through all layers except the last one
        for layer in list(self.layers)[:-1]:
            embed = layer(embed)
        return embed

class AdvantagePredictor(nn.Module):
    def __init__(self, dtype=None):
        super().__init__()
        self.dtype = dtype

        self.conv_layers = nn.Sequential(
            nn.Conv2d(in_channels=4, out_channels=32, kernel_size=3, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=2, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((2, 2))
        )
        self.fc_layers = nn.Sequential(
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(512, 128),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(128, 1)
        )
        self._initialize_weights()

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
                nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)

    def forward(self, x):
        x = self.conv_layers(x)
        x = x.view(x.size(0), -1)
        x = self.fc_layers(x)
        x = torch.sigmoid(x)
        return x.view(-1)*10

class AestheticScorerDiff(torch.nn.Module):
    def __init__(self, dtype):
        super().__init__()
        self.clip = CLIPModel.from_pretrained("openai/clip-vit-large-patch14")
        self.mlp = MLPDiff()
        state_dict = torch.load(ASSETS_PATH.joinpath("sac+logos+ava1-l14-linearMSE.pth"))
        self.mlp.load_state_dict(state_dict)
        self.dtype = dtype
        self.eval()

    def __call__(self, images):
        device = next(self.parameters()).device
        embed = self.clip.get_image_features(pixel_values=images)
        embed = embed / torch.linalg.vector_norm(embed, dim=-1, keepdim=True)
        return self.mlp(embed).squeeze(1), embed
    
    def generate_feats(self, images):
        device = next(self.parameters()).device
        embed = self.clip.get_image_features(pixel_values=images)
        embed = embed / torch.linalg.vector_norm(embed, dim=-1, keepdim=True)
        return embed

class offline_AestheticScorerDiff(torch.nn.Module):
    def __init__(self, dtype, config):
        super().__init__()
        self.clip = CLIPModel.from_pretrained("openai/clip-vit-large-patch14")
        self.dtype = dtype
        
        self.pessimism = config.train.pessimism

        if self.pessimism == 'none':
            self.mlp = MLPDiff_simple()
            state_dict = torch.load("./reward_aesthetic/models/baseline_simple.pth")
            self.mlp.load_state_dict(state_dict)
            
        elif self.pessimism == 'LCB':
            self.mlp = MLPDiff_simple()
            state_dict = torch.load("./reward_aesthetic/models/baseline_simple.pth")
            self.mlp.load_state_dict(state_dict)
            self.cov_mat = torch.load("./reward_aesthetic/models/AVA_covariances_simple.pt")
            
        elif self.pessimism == 'bootstrap':
            from reward_aesthetic.train_bootstrap import BootstrappedNetwork
            
            # The bootstrapped network consists 4 identical MLPDiff_simple() models
            self.model = BootstrappedNetwork(768)  # 768 is the dimension of CLIP embeddings
            state_dict = torch.load("./reward_aesthetic/models/bootstrap_simple.pth")
            self.model.load_state_dict(state_dict)
        
        else:
            raise NotImplementedError
        
        self.eval()

    def __call__(self, images, config):
        device = next(self.parameters()).device
        embed = self.clip.get_image_features(pixel_values=images)
        embed = embed / torch.linalg.vector_norm(embed, dim=-1, keepdim=True)
        
        if self.pessimism == 'none':
            return self.mlp(embed).squeeze(1), embed

        if self.pessimism == 'LCB':

            feats = self.mlp.forward_up_to_second_last(embed)  #(B,16)
            raw_rewards = self.mlp(embed)
            
            penalties = torch.zeros_like(raw_rewards).to(feats.device)

            for idx in range(raw_rewards.shape[0]):
                feat = feats[[idx,],].t()
                cov_matrix = self.cov_mat.to(feat.device) + config.train.psm_lambda * torch.eye(feat.shape[0]).to(feat.device)
                penalty = config.train.psm_alpha*torch.sqrt(torch.mm(feat.t(), torch.mm(torch.linalg.inv(cov_matrix), feat)))
                penalties[idx,] = penalty.squeeze(1)
            
            pessmistic_rewards = raw_rewards - torch.clamp(penalties,\
                            max=config.train.psm_clipping)
            
            return pessmistic_rewards.squeeze(1), embed
        
        elif self.pessimism == 'bootstrap': 
            outputs = self.model(embed)
            
            ##### Bootstrapping version 1: use a random head for output
            # rand_idx = random.randint(0, len(outputs) - 1) # use a random head for output
            # return outputs[rand_idx].squeeze(1)
            
            stacked_outputs = torch.cat(outputs, dim=1) #(B, N=10)
            
            ##### Bootstrapping version 2: use the softmin of all heads as output
            # tau = config.train.temp # Scale the outputs by the temperature parameter, compute the softmin along the second dimension (dim=1) to get a tensor shaped (B, N)
            # scaled_outputs = stacked_outputs / tau
            # weights = torch.nn.functional.softmin(scaled_outputs, dim=1)
            # weighted_sum = (weights * stacked_outputs).sum(dim=1)
            
            ##### Bootstrapping version 3: use the minimum value of the outputs
            pessmistic_rewards, _ = torch.min(stacked_outputs, dim=1, keepdim=True)
            
            return pessmistic_rewards.squeeze(1), embed


if __name__ == "__main__":
    scorer = AestheticScorerDiff(dtype=torch.float32).cuda()
    scorer.requires_grad_(False)
    for param in scorer.clip.parameters():
        assert not param.requires_grad
    for param in scorer.mlp.parameters():
        assert not param.requires_grad