'@Author:NavinKumarMNK'
import torch
from torch import nn
import sys
import os
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../../')))

import pytorch_lightning as pl
import models
from models.MalCoAtNet.MalCoAtNet import MalCoAtNet
from models.MalCoAtNet.RGB.MalwareImageDataset import MalBinImgDataModule
import utils.utils as utils
import wandb
import torchvision as tv
import ray_lightning as rl

class MalwareClassifer(pl.LightningModule):
    def __init__(self, hidden_rep, num_classes, weights_path,
                    in_channels):
        super(MalwareClassifer, self).__init__()
        self.model = MalCoAtNet(weights_path=weights_path,
                                 in_channels=in_channels)
        self.example_input_array = torch.randn(1, in_channels, 256, 256)
        self.example_output_array = torch.randn(1, num_classes)
        self.save_hyperparameters()
        self.classifier = nn.Sequential(
            nn.Linear(hidden_rep, hidden_rep),
            nn.BatchNorm1d(hidden_rep),
            nn.ReLU(),
            nn.Dropout(0.25),
            nn.Linear(hidden_rep, hidden_rep),
            nn.BatchNorm1d(hidden_rep),
            nn.ReLU(),
            nn.Dropout(0.25),
            nn.Linear(hidden_rep, 768),
            nn.BatchNorm1d(768),
            nn.ReLU(),
            nn.Linear(768, 256),
            nn.ReLU(),
            nn.Linear(256, 64),
            nn.ReLU(),
            nn.Linear(64, num_classes)
        )
        self.best_val_loss = None

        try:
            self.classifier = torch.load(utils.ROOT_PATH + '/weights/MalwareClassifierBB.pt')
        except FileNotFoundError:
            torch.save(self.classifier, utils.ROOT_PATH + '/weights/MalwareClassifierBB.pt')

    def forward(self, x):
        x = self.model(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x
    
    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = nn.CrossEntropyLoss()(logits, y)
        acc = (logits.argmax(dim=1) == y).float().mean()
        self.log('train_loss', loss)
        self.log('train_acc', acc, prog_bar=True)
        return {'loss': loss, 'acc': acc}

    def training_epoch_end(self, outputs):
        avg_loss = torch.stack([x['loss'] for x in outputs]).mean()
        avg_acc = torch.stack([x['acc'] for x in outputs]).mean()
        self.log('train/acc', avg_acc)
        self.log('train/loss', avg_loss)
        

    def validation_step(self, batch, batch_idx):
        print("hello")
        x, y = batch
        logits = self(x)
        loss = nn.CrossEntropyLoss()(logits, y)
        acc = (logits.argmax(dim=1) == y).float().mean()
        self.log('val_loss', loss, prog_bar=True)
        self.log('val_acc', acc, prog_bar=True)
        return {'val_loss': loss, 'val_acc': acc}

    def validation_epoch_end(self, outputs):
        print(outputs[0])
        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
        avg_acc = torch.stack([x['val_acc'] for x in outputs]).mean()
        self.log('val/loss', avg_loss, prog_bar=True)
        self.log('val/acc', avg_acc, prog_bar=True)
        if(self.best_val_loss) == None:
            self.best_val_loss = avg_loss
            #self.save_model()
        elif (avg_loss < self.best_val_loss):
            self.best_val_loss = avg_loss
            self.save_model()
        
    def test_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = nn.CrossEntropyLoss()(logits, y)
        acc = (logits.argmax(dim=1) == y).float().mean()
        self.log('test_loss', loss)
        self.log('test_acc', acc, prog_bar=True)
        return {'test_loss': loss}
    
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=0.001)
        lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
        return {
            'optimizer': optimizer, 
            'lr_scheduler': lr_scheduler
        }

    def save_model(self):
        self.model.save_model()
        torch.save(self.classifier, utils.ROOT_PATH + '/weights/MalwareClassifierBB.pt')
        artifact = wandb.Artifact('MalwareClassifierBB.cpkt', type='model')

if __name__ == '__main__' :
    wandb.init()
    from pytorch_lightning.loggers import WandbLogger
    logger = WandbLogger(project='Malware-Analysis', name='Malware Classifer')

    
    from pytorch_lightning import Trainer
    
    #import ray
    # #ray.init(runtime_env={"working_dir": utils.ROOT_PATH},
    #         )

    dataset_params = utils.config_parse('MALBINIMG_BB_DATASET')  
    dataset = MalBinImgDataModule(**dataset_params)
    dataset.setup() 
    print("dataset complete")

    from pytorch_lightning.callbacks import ModelSummary
    from pytorch_lightning.callbacks.progress import TQDMProgressBar
    from pytorch_lightning.callbacks import ModelCheckpoint
    from pytorch_lightning.callbacks import EarlyStopping
    from pytorch_lightning.callbacks.device_stats_monitor import DeviceStatsMonitor
    
    early_stopping = EarlyStopping(monitor='val/loss', patience=5)
    device_monitor = DeviceStatsMonitor()
    checkpoint_callback = ModelCheckpoint(dirpath=utils.ROOT_PATH + '/weights/checkpoints/malcoatnet/',
                                            monitor="val/loss", mode='min', every_n_train_steps=100, save_top_k=5, save_last=True)
    model_summary = ModelSummary(max_depth=5)
    refresh_rate = TQDMProgressBar(refresh_rate=10)

    callbacks = [
        model_summary,
        refresh_rate,
        checkpoint_callback,
        early_stopping,
        device_monitor
    ]

    mal_clf = utils.config_parse('MALWARE_CLASSIFIER_BB')
    model = MalwareClassifer(**mal_clf)
    mal_clf_params = utils.config_parse('MALWARE_CLASSIFIER_TRAIN')


    dist_env_params = utils.config_parse('DISTRIBUTED_ENV')
    strategy = None
    if int(dist_env_params['horovod']) == 1:
        strategy = rl.HorovodRayStrategy(num_workers=dist_env_params['num_workers'],
                                        use_gpu=dist_env_params['use_gpu'])
    elif int(dist_env_params['model_parallel']) == 1:
        strategy = rl.RayShardedStrategy(num_workers=dist_env_params['num_workers'],
                                        use_gpu=dist_env_params['use_gpu'])
    elif int(dist_env_params['data_parallel']) == 1:
        strategy = rl.RayStrategy(num_workers=dist_env_params['num_workers'], 
                                        use_gpu=dist_env_params['use_gpu'])
    
    trainer = Trainer(**mal_clf_params, 
                    callbacks=callbacks,
                    logger=logger,
		    accelerator='gpu',
                    #strategy='deepspeed_stage_1',
		    #num_nodes=4,
            log_every_n_steps=5,
            #resume_from_checkpoint=utils.ROOT_PATH + '/weights/checkpoints/malcoatnet/last.ckpt',
                    )

    trainer.fit(model, dataset)
    model.save_model()

    #model.model.finalize()

    #trainer.test(model, dataset.test_dataloader())
    wandb.finish()

    
    '''
    # SVM Data Preparation
    # inference of model
    model.model.eval()
    features = []
    for batch in dataset.train_dataloader():
        y = model.model.feature_extractor(batch[0])
        y = y.view(y.size(0), -1)
        label = batch[1]
        # combine y and label
        y = torch.cat((y, label.unsqueeze(1)), dim=1)
        features.append(y)
    '''

