import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
torch.manual_seed(123)
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

#GPU
# device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
mse = nn.MSELoss()
#mse function without reduction
mse_none = nn.MSELoss(reduction='none')

class MyDataset(torch.utils.data.Dataset):
  def __init__(self, X, Y):
        'Initialization'
        self.X = X
        self.Y = Y

  def __len__(self):
        'Denotes the total number of samples'
        return len(self.X)

  def __getitem__(self, index):
        'Generates one sample of data'
        X = self.X[index]
        Y = self.Y[index]

        return X, Y

class MyDataset2(torch.utils.data.Dataset):
  def __init__(self, X, Y1,Y2):
        'Initialization'
        self.X = X
        self.Y1 = Y1
        self.Y2 = Y2

  def __len__(self):
        'Denotes the total number of samples'
        return len(self.X)

  def __getitem__(self, index):
        'Generates one sample of data'
        X = self.X[index]
        Y1 = self.Y1[index]
        Y2 = self.Y2[index]

        return X, Y1, Y2
    


def eval_batch(net, train_loader,func, device):
  metrics=[]
  for data in train_loader:
    x,y = data
    result = net(x.to(device))
    loss = func(result.to('cpu'),y)
    metrics.append(loss.item())
  return torch.mean(torch.tensor(metrics)).item()

def eval_loader(model,loader,device):
  '''
  Input:  (1) model: NN model
          (2) loader for testing data: should have real/imaginary parts of spectrums
          (3) gpu device to train on
  '''
  loss = []
  for data in loader:
    x, y1,y2 = data
    x = x.to(device)
    y1 = y1.to(device)
    y2 = y2.to(device)
    output,_ = model(x)
    preal = output.real*x
    pimag = output.imag*x

    loss.append(mse(torch.cat([preal,pimag]),torch.cat([y1,y2])).item())

  return torch.mean(torch.tensor(loss))

