from tqdm import tqdm
import os
from functools import partial
import pandas as pd
from functools import partial
from torch.utils.data import DataLoader
from torch.utils.data import Dataset
import torch

# QML
import pennylane as qml
from pennylane import numpy as np
from pennylane.optimize import NesterovMomentumOptimizer, QNGOptimizer
from pennylane.templates.embeddings import AmplitudeEmbedding,AngleEmbedding



#################################################################
# DEFINING CIRCUITS
#################################################################

def circuit1(params, embedding, n_layers, n_qubits, x):
    embedding(x)

    for layer in range(n_layers):
        W = params[layer]

        # Define Rotations
        for i in range(n_qubits):
            qml.Rot(W[i, 0], W[i, 1], W[i, 2], wires=i)

        # CNots
        if n_qubits > 1: # Da erro quando n_qubits=1 coz nao pode haver CNots TODO?
            for i in range(n_qubits):
                if i == n_qubits-1:
                    qml.CNOT(wires=[i, 0])
                else:
                    qml.CNOT(wires=[i, i+1])

    return qml.expval(qml.PauliZ(0))


def circuit2(params, embedding, n_layers, n_qubits, x):
    embedding(x)

    for layer in range(n_layers):
        W = params[layer]
        for i in range(n_qubits):
            qml.Rot(W[i, 0], W[i, 1], W[i, 2], wires=i)

        if n_qubits > 1: # Da erro quando n_qubits=1 coz nao pode haver CNots TODO?
            for i in range(n_qubits):
                for j in range(i+1,n_qubits):
                    qml.CNOT(wires=[i, j])

    return qml.expval(qml.PauliZ(0))


#################################################################
# DEFINING DATALOADERS
#################################################################

class BankData(Dataset):
    """
    This will load the data
    """
    def __init__(self, category, standardization='ML'):

        # Load data_banknote_authentication.txt
        # https://archive.ics.uci.edu/ml/machine-learning-databases/00267/data_banknote_authentication.txt
        # 1 Variance of Wavelet Transformed image (continuous).
        # 2 Skewness of Wavelet Transformed image (continuous).
        # 3 Kurtosis of Wavelet Transformed image (continuous).
        # 4 Entropy of image (continuous).
        # 5 label (0 for authentic, 1 for inauthentic).
        if 'data_banknote_authentication.txt' not in os.listdir():
            os.system('wget https://archive.ics.uci.edu/ml/machine-learning-databases/00267/data_banknote_authentication.txt')


        self.data = pd.DataFrame(np.loadtxt('data_banknote_authentication.txt', delimiter=','))
        self.data.columns = ['var', 'skew', 'kurt', 'ent', 'label']

        # Sanity checks
        assert category in {'train', 'validation', 'test', 'all'},   "Invalid category!"


        # Standardization
        _features = self.data.columns[:-1] # Everything except label

        if standardization == 'ML':
            self.data[_features] = (self.data[_features]-self.data[_features].mean())/self.data[_features].std()

        elif standardization == 'AngleEmbedding':
            # Put the data between -pi and pi
            self.data[_features] = self.data[_features] % (2*np.pi)
            pass
        elif standardization == 'AmplitudeEmbedding':
            pass
        else:
            # Error
            raise ValueError("Invalid standardization!")


        # This will equally devide the dataset into
        # train, validation and test
        train, validation, test = np.split(self.data, [int(len(self.data)*(1/3)), int(len(self.data)*(2/3))])

        if category == "train":
            self.data = train
        elif category == "validation":
            self.data = validation
        elif category == "test":
            self.data = test
        elif category == "all":
            pass
        else:
            # error
            raise ValueError("Invalid category!")

        del train, validation, test

        self.label = self.data["label"]
        self.data.drop(columns=["label"], inplace=True)

        self.n_samples = self.data.shape[0]

    def get_columns(self):
        return self.data.columns

    def __getitem__(self, index):
        # Data, Label
        return  torch.tensor(self.data.iloc[index]),      \
                torch.tensor(self.label.iloc[index])

    def __len__(self):
        return self.n_samples


#################################################################
# TRAINING LOOP
#################################################################


class Trainer():
    def __init__(self, parameteres):
        # Get parameters
        self.max_epochs = parameteres['max_epochs']
        self.batch_size = parameteres['batch_size']
        self.circuit_name = parameteres['circuit_name']
        self.n_features = parameteres['n_features']
        self.embedding_method = parameteres['embedding_method']
        self.n_layers = parameteres['n_layers']

        # This assumes I'm using angle embedding for demonstration purposes
        self.n_qubits = self.n_features
        self.embedding = partial(globals()[self.embedding_method], wires=range(self.n_qubits), rotation='Z')
        self.normalization = self.embedding_method

        # Defining device
        self.dev = qml.device(parameteres['device'], wires=self.n_qubits)
        self.opt = parameteres['optimizer']

        # Parameter Initialization
        self.weights = 0.01 * np.random.randn(self.n_layers, self.n_qubits, 3, requires_grad=True)
        self.bias = np.array(0.0, requires_grad=True)

        # Define circuit
        self.circuit_ = qml.QNode(globals()[self.circuit_name], self.dev)

        # Load data
        self.load_data()


    def circuit(self, weights, x):
        return self.circuit_(weights, self.embedding, self.n_layers, self.n_qubits, x)

    def classifier(self, weights, bias, x):
        return self.circuit(weights, x) + bias

    def cost(self, weights, bias, X, Y): # weights, bias, X_batch, Y_batch
        predictions = np.array([np.sign(self.classifier(weights, bias, x)) for x in X])
        loss = np.mean((predictions - Y)**2)

        return loss

    def train(self):
        for epoch in range(self.max_epochs):
            self.epoch = epoch
            # Create batch iterator
            dataloader = DataLoader(dataset=self.train_dataset, batch_size=self.batch_size, num_workers=os.cpu_count())

            for i, pack in tqdm(enumerate(iter(dataloader)), total=len(self.train_dataset)/self.batch_size, desc=f"Training epoch {epoch+1}/{self.max_epochs} - "):
                # Get current Step
                self.global_step = epoch * (len(self.train_dataset)/self.batch_size) + i
                loss = self.training_step(i, pack)

            # Print to terminal
            print("Epoch: {:5d} | Step: {:5d} | Cost: {:0.7f} ".format(
                        epoch, i, loss
                    ))



    def training_step(self, i, pack):
        x_batch, y_batch = pack

        x_batch = np.array(x_batch, requires_grad=False)
        y_batch = np.array(y_batch, requires_grad=False)

        # Replace 0 with -1 on y due to the architecture of the circuit TODO: Maybe change loss?
        y_batch[y_batch == 0] = -1

        # Other things I've tried
        print("params before\n", self.weights, "\n", self.bias)
        pack, loss = self.opt.step_and_cost(self.cost, self.weights, self.bias, x_batch, y_batch)
        self.weights, self.bias, _, _ = pack
        # Also tried this (again) and only get a gradient on the bias
        #grad = self.opt.compute_grad(self.cost, (self.weights, self.bias), {'X':x_batch, 'Y':y_batch})
        #self.weights, self.bias = self.opt.apply_grad(grad, (self.weights, self.bias))
        print("params after\n", self.weights, "\n", self.bias)
        return loss


    def load_data(self):
        # Datasets initialization / Load data
        print("[Info] Loading data...")
        self.train_dataset = BankData('train')
        self.val_dataset =   BankData('validation')




dummy_params = {
    'device': 'default.qubit',
    'batch_size': 10,
    'max_epochs': 4,
    'n_datapoints': 10000,
    'patience': 20,
    'check_val_every_n_epoch': 2,
    'optimizer': NesterovMomentumOptimizer(100),
    'n_layers': 1,
    'embedding_method': 'AngleEmbedding',
    'circuit_name':'circuit1',
    'n_features': 4,
}

machine = Trainer(dummy_params)
machine.train()