import pandas as pd

import numpy as np

import sklearn.metrics as skm

import matplotlib.pyplot as plt

from sklearn.model_selection import train_test_split

from sklearn.preprocessing import LabelEncoder, StandardScaler

import torch

from torch import nn

from torch import optim

from torch.nn import functional as F

from torch.optim.lr_scheduler import _LRScheduler

from torch.utils.data import TensorDataset, DataLoader

from tqdm import tqdm, trange

from torchsummary import summary

import getopt, sys

from tsai.imports import *

from tsai.models.layers import *

from tsai.models.utils import *

def noop(x):

    pass

def shortcut(c_in, c_out):

    return nn.Sequential(*[nn.Conv1d(c_in, c_out, kernel_size=1), 

                           nn.BatchNorm1d(c_out)])

def convert_sig(x):

    return(0 if x<0.5 else 1)

class InceptionModule(Module):

    def __init__(self, ni, nf, ks=40, bottleneck=True):

        ks = [ks // (2**i) for i in range(3)]

        ks = [k if k % 2 != 0 else k - 1 for k in ks]  # ensure odd ks

        bottleneck = bottleneck if ni > 0 else False

        self.bottleneck = Conv1d(ni, nf, 1, bias=False) if bottleneck else noop

        self.convs = nn.ModuleList([Conv1d(nf if bottleneck else ni, nf, k, bias=False) for k in ks])

        self.maxconvpool = nn.Sequential(*[nn.MaxPool1d(3, stride=1, padding=1), Conv1d(ni, nf, 1, bias=False)])

        self.concat = Concat()

        self.bn = BN1d(nf * 4)

        self.act = nn.ReLU()

    def forward(self, x):

        input_tensor = x

        x = self.bottleneck(input_tensor)

        x = self.concat([l(x) for l in self.convs] + [self.maxconvpool(input_tensor)])

        return self.act(self.bn(x))

@delegates(InceptionModule.__init__)

class InceptionBlock(Module):

    def __init__(self, ni, nf=32, residual=True, depth = 6, **kwargs):

        self.residual, self.depth = residual, depth

        self.inception, self.shortcut = nn.ModuleList(), nn.ModuleList()

        for d in range(depth):

            self.inception.append(InceptionModule(ni if d == 0 else nf * 4, nf, **kwargs))

            if self.residual and d % 3 == 2: 

                n_in, n_out = ni if d == 2 else nf * 4, nf * 4

                self.shortcut.append(BN1d(n_in) if n_in == n_out else ConvBlock(n_in, n_out, 1, act=None))

        self.add = Add()

        self.act = nn.ReLU()

    def forward(self, x):

        res = x

        for d, l in enumerate(range(self.depth)):

            x = self.inception[d](x)

            if self.residual and d % 3 == 2: res = x = self.act(self.add(x, self.shortcut[d//3](res)))

        return x

@delegates(InceptionModule.__init__)

class InceptionTime(Module):

    def __init__(self, c_in, c_out, nf=32, nb_filters=None, **kwargs):

        nf = ifnone(nf, nb_filters) # for compatibility

        self.inceptionblock = InceptionBlock(c_in, nf, **kwargs)

        self.gap = GAP1d(1)

        self.fc = nn.Linear(nf * 4, c_out)

        self.sig = nn.Sigmoid()

    def forward(self, x):

        x = self.inceptionblock(x)

        x = self.gap(x)

        x = self.sig(self.fc(x))

        return x

def train_model(train_dl: DataLoader,

          test_dl: DataLoader,

          device: str,

          model: nn.Module,

          epochs: int,

          learning_rate: float,

          Save: bool):

    optimiser = optim.Adam(model.parameters() ,lr=learning_rate)

    history = []

    loss_history = []

    acc_history = []

    epoch_bar = trange(epochs)

    for epoch in epoch_bar:

        epoch_loss = 0

        model.train(mode = True)

        for batch, data in enumerate(train_dl):

            x,y = data

            x,y = x.to(device),y.to(device)

            #computation graph (forward prop ->compute loss ->back prop ->update weights)

            optimiser.zero_grad()

            out = model(x)

            y = torch.Tensor(y.cpu().detach().numpy()).view(y.shape[0], 1).to(device)

            loss = loss_func(out, y)

            epoch_loss += loss.item()

            loss.backward()

            optimiser.step()

        loss_history.append(epoch_loss)

        print ("Train Loss: ",epoch_loss/len(train_dl))

        #Validation

        running_loss = 0

        running_acc  = 0

        running_far = 0

        model.eval()

        for batch, data in enumerate(test_dl):

            x, y = data

            x, y = x.to(device),y.to(device)

            out = model(x)

            convert_soft = np.vectorize(convert_sig)

            out1 = torch.Tensor(convert_soft(out.cpu().detach().numpy())).view(-1).to(device)

            test_acc = (out1 == y.view(-1)).cpu().detach().numpy().sum()/len(y)

            y = torch.Tensor(y.cpu().detach().numpy()).view(y.shape[0], 1).to(device)            

            test_loss = loss_func(out,y).item()

            running_acc  += test_acc

            running_loss += test_loss

        test_size = len(test_dl)

        test_acc = running_acc/(batch+1)

        test_loss = running_loss/(batch+1)

        epoch_bar.set_description('acc={0:.2f}%\tBCE={1:.4f}%'

                                  .format(test_acc, test_loss))

    return model,history 

if __name__ == "__main__":

    arch = InceptionTime(1, 1)

    df = pd.read_csv('finaldfs/ecgfiltered30sec.csv', index_col = 0)

    train = df.groupby('infant_no').apply(lambda group : group[group['brady_no'] <= (group['brady_no'].max())*0.7]).copy()

    test = df.groupby('infant_no').apply(lambda group : group[group['brady_no'] > (group['brady_no'].max())*0.7]).copy()

    x_train = train[train.columns[4:-2]]

    y_train = train['brady']

    x_test = test[test.columns[4:-2]]

    y_test = test['brady']

    x_train = np.expand_dims(x_train, axis = 1)

    x_test = np.expand_dims(x_test, axis = 1)

    x_train = torch.Tensor(x_train)

    x_test = torch.Tensor(x_test)

    y_train = torch.Tensor(y_train.to_numpy())

    y_test = torch.Tensor(y_test.to_numpy())

    train_ds = TensorDataset(x_train, y_train)

    test_ds = TensorDataset(x_test, y_test)

    train_dl = DataLoader(train_ds, batch_size = 10, shuffle = True)

    test_dl = DataLoader(test_ds, batch_size = 10, shuffle = True)

    device = torch.device('cuda:0')

    train_size = x_train.shape[0]

    test_size = x_test.shape[0]

    time_steps = x_train.shape[-1]

    num_classes = 1

    learning_rate = 1e-6

    drop = 0.2

    epochs = 100

    loss_func = nn.BCELoss()

    model = arch

    model = model.to(device)

    train_model(train_dl,

          test_dl,

          device,

          model,

          epochs,

          learning_rate, False)