import torch
import os
from tqdm import tqdm
from torch import Tensor
import torch.nn.functional as F
import numpy as np
from dataset_v2 import TrainDataset
from model import UNet, MultiLoss
from torch.utils.data import DataLoader
#import pandas as pd
#import torchsummary
import torch.nn as nn
import torchaudio
from util import tf2pytorch

params = {
    ### Dataset ###
    'margin': 0.5,
    'chunk_duration': 20.0,
    'sample_rate': 44100,
    'frame_length': 4096,
    'frame_step': 1024,
    'T': 512,
    'F': 1024,
    'n_chunks_per_song': 15,
    'train_manifest': 'mass_dataset\\manifest.csv', # Manifest generated by preprocess.py

    ### Train ###
    'epochs': 100000,
    'batch_size': 4,
    'optimizer': 'adam',
    'loss': 'l1',
    'momentum': 0.9,
    'dampening': 0,
    'lr': 5e-5,
    'lr_decay': 0,
    'wd': 0.00001,
    'model_dir': './model/',
    'final_dir': './final_model/',
    'load_optimizer': True,
    'start': None,
    'load_model': 'pytorch', # 'pytorch' / 'tensorflow' / None -> new model
    'resume': ['./final_model/net_vocal.pth', './final_model/net_instrumental.pth'],
    'checkpoint_path': './2stems/model',
    'num_instruments': ['vocal', 'instrumental'],
    'seed': 123456
}

def load_ckpt(model, ckpt):
        state_dict = model.state_dict()
        for k, v in ckpt.items():
            if k in state_dict:
                target_shape = state_dict[k].shape
                assert target_shape == v.shape
                state_dict.update({k: torch.from_numpy(v)})
            else:
                print('Ignore ', k)

        model.load_state_dict(state_dict)
        return model


def train(epoch, model_list, multi_loss, criterion, optimizer, train_loader, params):
    for model in model_list:
        model.train()

    sum_loss, sum_samples = 0, 0
    progress_bar = tqdm(enumerate(train_loader))
    for batch_idx, (mix_stft_mag, vocal_stft_mag, instru_stft_mag) in progress_bar:
        sum_samples += len(mix_stft_mag)

        mix_stft_mag = mix_stft_mag.transpose(2, 3)
        mix_stft_mag = mix_stft_mag.to(device)
        
        separate_stft_mag = []
        vocal_stft_mag = vocal_stft_mag.transpose(2, 3)
        instru_stft_mag = instru_stft_mag.transpose(2, 3)
        vocal_stft_mag = vocal_stft_mag.to(device)
        instru_stft_mag = instru_stft_mag.to(device)
        separate_stft_mag.append(vocal_stft_mag)
        separate_stft_mag.append(instru_stft_mag)

        loss = multi_loss(mix_stft_mag, separate_stft_mag)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        sum_loss += loss.item() * len(mix_stft_mag)
        progress_bar.set_description(
            'Train Epoch: {:3d} [{:4d}/{:4d} ({:3.3f}%)] Loss: {:.4f}'.format(
                epoch, batch_idx + 1, len(train_loader),
                100. * (batch_idx + 1) / len(train_loader),
                sum_loss / sum_samples))

    if epoch%200==0:
        for i in range(len(params['num_instruments'])):
            torch.save({'epoch': epoch, 'state_dict': model_list[i].state_dict(),
                    'optimizer': optimizer.state_dict()},
                   '{}/net_{}_{}.pth'.format(params['model_dir'], params['num_instruments'][i], epoch))
            torch.save({'epoch': epoch, 'state_dict': model_list[i].state_dict(),
                    'optimizer': optimizer.state_dict()},
                   '{}/net_{}.pth'.format(params['final_dir'], params['num_instruments'][i]))    




def main(params):
    torch.manual_seed(params['seed'])
    print(torch.version.cuda)
    train_dataset = TrainDataset(params)
    n_chunks = train_dataset.count
    print('Num of Chunks: {}'.format(n_chunks))

    model_list = nn.ModuleList()
    start = 1

    if params['load_model'] == 'pytorch':
        print('=> loading checkpoint from PyTorch {}'.format(params['resume']))
        for i in range(len(params['num_instruments'])):
            checkpoint = torch.load(params['resume'][i])
            net = UNet()
            if params['start'] is not None:
                start = int(params['start'])
            else:
                start = checkpoint['epoch'] + 1
            #if params['load_optimizer']:
            #    optimizer.load_state_dict(checkpoint['optimizer'])
            net.load_state_dict(checkpoint['state_dict'])
            net.to(device)
            model_list.append(net)
    elif params['load_model'] == 'tensorflow':
        print('=> loading checkpoint from Tensorflow {}'.format(params['checkpoint_path']))
        ckpts = tf2pytorch(params['checkpoint_path'], params['num_instruments'])
        for i in range(len(params['num_instruments'])):
            print('Loading model for instrument {}'.format(i))
            net = UNet()
            ckpt = ckpts[i]
            net = load_ckpt(net, ckpt)
            net.to(device)
            model_list.append(net)
    else:
        print('=> no checkpoint found at {}'.format(params['resume']))
        for i in range(len(params['num_instruments'])):
            net = UNet()
            net.to(device)
            model_list.append(net)


    if params['loss'] == 'l1':
        criterion = nn.L1Loss()
    else:
        criterion = nn.MSELoss()

    multi_loss = MultiLoss(model_list, criterion, params)

    if params['optimizer'] == 'sgd':
        optimizer = torch.optim.SGD(multi_loss.parameters(), lr=params['lr'], momentum=params['momentum'], dampening=params['dampening'], weight_decay=params['wd'])
    elif params['optimizer'] == 'adagrad':
        optimizer = torch.optim.Adagrad(multi_loss.parameters(), lr=params['lr'], lr_decay=params['momentum'], weight_decay=params['wd'])
    else:
        optimizer = torch.optim.Adam(multi_loss.parameters(), lr=params['lr'], weight_decay=params['wd'])


    train_loader = DataLoader(train_dataset, batch_size=params['batch_size'], shuffle=True,
                              num_workers=0, pin_memory=True)

    for epoch in range(start, params['epochs']):
        train(epoch, model_list, multi_loss, criterion, optimizer, train_loader, params)
  
device = torch.device('cuda')

# 'model_dir': './model/',
# 'final_dir': './final_model/',

os.makedirs(params['model_dir'], exist_ok=True)
os.makedirs(params['final_dir'], exist_ok=True)

main(params)

