from classifier import *
import numpy as np
from generator import *
from discriminator import *
from config import *
from dataset import *
from model import *

model = Classifier()
model.to(DEVICE)

optim = torch.optim.Adam(model.parameters(),
                            lr = LR_C, betas = (0.9, 0.999))

criterion = nn.NLLLoss()


def train_epoch(epoch, print_freq=None):
    model.train()
    train_loss = 0

    for i, (img, trg) in enumerate(train_dataloader):
        model.zero_grad()

        img, trg = img.to(DEVICE), trg.to(DEVICE)
        out = model(img)
        loss = criterion(out, trg)
        loss.backward()
        optim.step()
        train_loss += loss

        if print_freq and (i+1) % print_freq == 0:
            
            print(f"BATCH: {i+1}/{len(train_dataloader)}:\t"
                 f"Training Loss: {train_loss / (i+1):.3f}"
                 ) 
    
    train_loss /= len(train_dataloader)
    print(f"Train loss: {train_loss}")
    
    return train_loss.item()


def test_epoch(epoch, dataloader):
    model.eval()
    test_loss = 0

    with torch.no_grad():
        for i, (img, trg) in enumerate(dataloader):
            
            img, trg = img.to(DEVICE), trg.to(DEVICE)
            out = model(img)
            loss = criterion(out, trg)
            test_loss += loss
    
    test_loss /= len(test_dataloader)
    print(f"Test loss: {test_loss}")
    
    return test_loss.item()


if TRAIN_CLASSIFIER :
    print('Training Classifier model.......')

    train_losses = []
    valid_losses = []

    mini = np.inf
    for epoch in range(C_EPOCHS):
        print(f"EPOCH: {epoch+1}", end='\n')
        
        train_loss = train_epoch(epoch)
        valid_loss = test_epoch(epoch, dataloader = valid_dataloader)
        train_losses.append(train_loss)
        valid_losses.append(valid_loss)

        if valid_loss < mini:
            mini = valid_loss
            torch.save(model.state_dict(), CPKT+f"CLASSIFIER_BEST_WEIGHTS.pth")

    plt.plot(train_losses, label='Train_loss')
    plt.plot(valid_losses, label='Valid_loss')
    plt.legend()
    plt.savefig("figures/classifier_Loss_curves.png")

    print('TEST DATA LOSS:')
    test_loss = test_epoch(0, test_dataloader)


if PRED_ON_GAN_OUTPUT:

    generator = Generator(latent_dim = LATENT_DIM, g_channels = G_CHANNELS,
                            batch_norm=BATCH_NORM, g_spectral_norm=G_SPECTRAL_NORM)

    # Loading generator model after last epoch
    generator.load_state_dict(torch.load(CPKT+f"GANS_EPOCH_{EPOCHS}.pth"))

    discriminator = Discriminator(d_spectral_norm = D_SPECTRAL_NORM)

    generator = generator.to(DEVICE)
    discriminator = discriminator.to(DEVICE)

    model = Model(
        latent_dim = LATENT_DIM,
        dataloader = train_dataloader,
        generator = generator, 
        discriminator = discriminator,
        batch_size = BATCH_SIZE,
        lr_g = LR_G,
        lr_d = LR_D
        )

    classifier = Classifier()
    classifier.to(DEVICE)
    classifier.load_state_dict(torch.load(CPKT+f"CLASSIFIER_BEST_WEIGHTS.pth"))

    latent_vecs = GET_NOISE(NUM_SAMPLES)

    with torch.no_grad():
        generated_images = generator(latent_vecs)
        generated_images = generated_images.to(DEVICE)
        prediction_on_GANs = classifier(generated_images)
        
    predictions = np.argmax(prediction_on_GANs.cpu().numpy(), axis=1)


    if PLOT_20:
        '''
            Plots 20 images generated by GANs for analysis
        '''
        generated_images = generated_images * -1
        fig, axes = plt.subplots(5,4,figsize=(10,10))

        for i in range(5):
            for j in range(4):
                axes[i][j].imshow(generated_images[i*4+j].cpu().numpy().transpose(1,2,0), cmap='gray')
                axes[i][j].axis('off')
                axes[i][j].set_title(f"Predicted: {predictions[i*4+j]}")
                
        plt.savefig('figures/Analysis_20.png')
    
    if PLOT_DISTRIBUTION:
        '''
            Plots distribution of generated samples over classes
        '''
        plt.hist(predictions)
        plt.xticks(np.arange(10))
        plt.ylabel('Number of images')
        plt.xlabel('Class')
        plt.savefig("figures/Distribution_of_GAN_images.png")
        plt.close()
