from loss import Loss
from feature_maps import LOGGER, FeatureMaps
import torch.optim as optim
from torchvision.utils import save_image
import matplotlib.pyplot as plt
plt.ion()
class Optimizer:
    @staticmethod
    def gradient_descent(content_img, style_img, content_img_clone):
        '''
        Fuction to apply gradient descent on the content image
        @params
        content_img: Original Image
        style_img: Styling Image
        content_img_clone: Copy of Original Image
        Author: @gaurangathavale
        '''
        LOGGER.info("Running gradient descent with the following parameters")
        epoch = 5000
        learning_rate = 0.1
        alpha = 1
        beta = 0
        LOGGER.info(f"{epoch},{learning_rate},{alpha},{beta}")

        optimizer=optim.Adam([content_img_clone],lr=learning_rate)
        LOGGER.info("Optimizer = " + str(optimizer))
        #fig = plt.figure()
        #ax = fig.add_subplot(111)
        feature_maps = FeatureMaps()
        for e in range(epoch):
    
            content_fmaps = feature_maps.get_fmaps(content_img)
            style_fmaps = feature_maps.get_fmaps(style_img)
            content_generated_fmaps = feature_maps.get_fmaps(content_img_clone)
            
            total_loss,total_cont_loss,total_style_loss = Loss.total_loss(alpha, beta, content_fmaps, style_fmaps, content_generated_fmaps)
            
            # clears x.grad for every parameter x in the optimizer. 
            # It’s important to call this before total_loss.backward(), otherwise it will accumulate the gradients from multiple passes.
            optimizer.zero_grad() 
            
            # total_loss.backward() computes dtotal_loss/dx for every parameter x which has requires_grad=True
            total_loss.backward()
            
            # Optimization Step / Update Rule
            optimizer.step()
            #plt.clf()
            #plt.plot(content_img_clone)
            if(e%100==0):
                learning_rate = max(learning_rate/2,0.001)
            if(e%10 == 0):
                LOGGER.info(f"Epoch = {e} Learning Rate = {learning_rate} Total Loss = {total_loss} Style Loss = {total_style_loss} Content Loss = {total_cont_loss}")
                name = "styled_mix/styled_" + str(e) +".png"
                save_image(content_img_clone,name)
