# -*- coding: utf-8 -*-
"""
Created on Wed Apr 21 22:08:30 2021

@author: arash
"""

import torch, glob, random
import torch.nn as nn
import torch.nn.functional as F
from modules import Net, LineNet
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
#from torchvision import transforms
from data_loader import my_collate, CustomIterableDataset, LineDataset, LineDatasetBalanced
import random, pickle
import numpy as np

import torch.multiprocessing as mp

from torch.nn.parallel import DistributedDataParallel as DDP




def trainingNetwork(dataset):
    
    #N=17481
    N=20086
    
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model = Net(input_dim=20086).cuda()
    model=nn.DataParallel(model, device_ids=[0,1], output_device=[0])
    model.to(device)
    #print('model')
    
    #loss_function = nn.MSELoss()
    loss_function = nn.BCEWithLogitsLoss().cuda()
    #loss_function = nn.BCELoss().cuda()
    
    #optimizer = optim.SGD(model.parameters(), lr=0.01, momentum =0.99)
    optimizer = optim.Adam(model.parameters(), lr=0.01)
    loss_tot=[]
    
    temp_loss=0    
    for epoch in range(100):
        if epoch % 1 ==0:
            print('Epoch ',epoch, '---> Loss ', temp_loss)
            
        dataset.randdata()
        X=DataLoader(dataset, batch_size = 256, collate_fn=my_collate, num_workers=8, shuffle=True)
        
        temp_loss=0
        
        for local_batch, local_labels in X:
            
            local_batch=torch.permute(local_batch,(1,0,2))
        
            #print(local_batch.size())
            #print(local_labels.size())
            optimizer.zero_grad()
            tar = model(local_batch.float().cuda())#, torch.Tensor(local_labels).float().cuda())
            #print(tar[:,0])
            #print(local_labels)
            #print(local_labels.size(), tar[:,0].size())
            loss = loss_function(tar[:,0].cuda(), local_labels.float().cuda()) #
            temp_loss+=loss.detach().item()
            loss.backward()
            optimizer.step()
         
        name='model3N' + str(epoch) + '.pt'    
        torch.save(model, name)    
        loss_tot.append(temp_loss)   
        pickle.dump(loss_tot, open('loss3N.p','wb'))

        
    return model, loss_tot      






def trainingLineNetwork():
    
    #N=17481
    N=20086
    
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")   
    
    model = LineNet(input_dim=20086).cuda()
    model=nn.DataParallel(model, device_ids=[0,1,2,3], output_device=[0])
    model.to(device)
    
    files=glob.glob('training_lined*.p')
    
    
    loss_function = nn.BCEWithLogitsLoss().cuda()
    
    
    #optimizer = optim.SGD(model.parameters(), lr=0.01, momentum =0.99)
    optimizer = optim.Adam(model.parameters(), lr=0.01)
    loss_tot=[]
    
    temp_loss=0    
    for epoch in range(3):
        if epoch % 1 ==0:
            print('Epoch ',epoch, '---> Loss ', temp_loss)
            
        #dataset.randdata()
        
        random.shuffle(files)
        
        for i in range(len(files)):
            print('files   ',i)
            dataset=LineDatasetBalanced(files[i])
            
            X=DataLoader(dataset, batch_size = 256, shuffle=True) #, num_workers=2
            
            temp_loss=0
            
            for local_batch, context_batch, local_labels in X:
                
                #local_batch=torch.permute(local_batch,(1,0,2))
            
                #print(local_batch.size())
                #print(local_labels.size())
                optimizer.zero_grad()
                
                #print(local_batch.size())            
                
                tar = model(local_batch.float().cuda(), context_batch.float().cuda())#, torch.Tensor(local_labels).float().cuda())
                #print(tar[:,0])
                #print(local_labels)
                #print(local_labels.size(), tar[:,0].size())
                loss = loss_function(tar[:,0].cuda(), local_labels[:,0].float().cuda()) #
                temp_loss+=loss.detach().item()
                loss.backward()
                optimizer.step()
         
        #name='modelLine' + str(epoch) + '.pt'    
        #torch.save(model, name)    
        loss_tot.append(temp_loss)   
        #pickle.dump(loss_tot, open('lossLine.p','wb'))

        
    return model, loss_tot      












