#!/usr/bin/env python
# -*- coding: utf-8 -*-

from dataclasses import dataclass
from flask import Flask, request
from threading import Thread
from multiprocessing import Queue
import traceback
import pickle
from dataclasses import dataclass
from num2words import num2words
from torch.optim import AdamW
from torch.utils.data import DataLoader, Dataset
from torchtext.data import get_tokenizer
from tqdm import tqdm
import torch.nn.functional as F
from typing import Dict, List, Set, Tuple
import splogger as log
from torch.utils.data import DataLoader, Dataset
from torchvision.datasets.cifar import CIFAR10
from tqdm import tqdm
from typing import List, NamedTuple, Tuple
from flask import jsonify
import numpy as np
import os
from flask_cors import CORS
import re
import time
import torch
import torchvision.transforms as T
from multiprocessing import Lock
import os
from werkzeug.utils import secure_filename

#! Put this in config ? IDK
α                   = .33
MAX_SEQ             = 50
ROOT                = "./dataset/full.txt"
BATCH_SIZE          = 1
NUM_WORKERS         = 4
RUNS                = 50
VOCAB               = "vocab.pkl"
FILE_UPLOAD_FOLDER  = './uploads'

# if not os.path.isfile(VOCAB):
#     hpb = HarryPotterBooks.from_file(ROOT, fraction=FRACTION)
#     with open(VOCAB, "wb") as f:
#         pickle.dump([
#             hpb.path,
#             hpb.raw,
#             hpb.tokens,
#             hpb.vocab,
#             hpb.word2id,
#             hpb.id2word,
#         ], f)


@dataclass
class HarryPotterBooks:
    path: str
    raw: str
    tokens: List[str]
    vocab: Set[str]
    word2id: Dict[str, int]
    id2word: Dict[int, str]

    @classmethod
    def from_file(cls, path: str, fraction: float = 1):
        with open(path, "r") as f:
            raw = f.read()
            raw = " ".join(raw.split())

        numbers = lambda t: num2words(t) if type(t) in [int, float] else t
        tokenizer = get_tokenizer("basic_english")
        tokens = tokenizer(raw)
        tokens = list(map(numbers, tokens))
        tokens = tokens[:int(len(tokens) * fraction)]

        vocab = set(tokens)

        word2id = {t: i for i, t in enumerate(vocab)}
        id2word = {i: t for i, t in enumerate(vocab)}
        
        return cls(path, raw, tokens, vocab, word2id, id2word)


class HarryPotterDataset(Dataset):
    def __init__(self, hpb: HarryPotterBooks, window: int) -> None:
        super(HarryPotterDataset, self).__init__()
        self.hpb = hpb
        self.window = window

    def __len__(self) -> int:
        return len(self.hpb.tokens) - (self.window + 1)

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]:
        data = self.hpb.tokens[idx:idx + self.window + 1]
        sentence_tokens, word_token = data[:-1], data[-1]

        sentence = torch.tensor([
            self.hpb.word2id[t] for t in sentence_tokens
        ], dtype=torch.long)
        word = self.hpb.word2id[word_token]
        
        return sentence, word


@dataclass
class Record():
    name:           str
    surname:        str
    perplexity:     float = 0.
    performance:    float = 0.
    memory:         float = 0.
    score_unscaled: float = 0.
    score:          float = 0.

    def update_score_unscaled(self, min_max: dict):
        # self.score_unscaled  = scale(self.accuracy, *min_max['accuracy']) + .5 * scale(self.performance, *min_max['performance']) + .5 * scale(self.memory, *min_max['memory'])
        self.score_unscaled = .9 * (1. - np.log(self.perplexity)/np.log(min_max['perplexity'][1])) + .001 * self.performance / min_max['performance'][1] + .001 * (1 - self.memory / min_max['memory'][1])
        # self.score_unscaled /= 3

    def scale_score(self, σ_unscaled):
        self.score = self.score_unscaled * (1 - α * σ_unscaled * self.score_unscaled)


class Entry:
    def __init__(self, path: str) -> None:
        self.path = path
        self.model = torch.jit.load(path).eval()

    @property
    def name(self) -> str:
        return self.path.split("/")[-1].split(".")[0].split("_")[1]

    @property
    def surname(self) -> str:
        return self.path.split("/")[-1].split(".")[0].split("_")[2]

    def accuracy(self, dataset: Dataset, loader: DataLoader) -> float:
        self.model = self.model.cpu()
        total_acc = 0.0
        for inputs, labels in loader:
            inputs, labels = inputs.cpu(), labels.cpu()
            outputs = self.model(inputs)
            total_acc += (
                (torch.argmax(outputs, 1) == labels).sum().item() /
                len(dataset)
            )
        return total_acc

    def perplexity(self, dataset: HarryPotterDataset, start: int = 50) -> float:
        n = 10
        self.model = self.model.cpu()
        prod = 1
        for k in range(start, start + n):
            sentence_ids, wid = dataset[k]

            _word_id = F.softmax(self.model(sentence_ids.unsqueeze(0).cpu()), dim=1)
            wk = _word_id.squeeze(0).cpu()[wid].item()

            prod *= 1/wk
        return prod ** 1/n

    def performance(
        self, input_size: Tuple[int, int, int], runs: int,
    ) -> float:
        self.model = self.model.cpu()
        times = np.zeros((runs, ))
        inputs = torch.randint(10, (1, input_size)).cpu()
        for i in range(runs):
            t0 = time.time()
            self.model(inputs)
            times[i] = time.time() - t0
        return 1 / np.median(times)

    def memory(self, input_size: Tuple[int, int, int]) -> float:
        self.model = self.model.cpu()
        size = 0.0
        for parameter in self.model.parameters():
            size += np.sum(parameter.size())
        return size * 32 * 1.192e-7

    def __call__(
        self,
        dataset: Dataset,
        loader: DataLoader,
        input_size: Tuple[int, int, int],
        runs: int,
    ) -> Record:
        return Record(
            self.name, self.surname,
            min(self.perplexity(dataset), 1e308),
            min(self.performance(input_size, runs), 60.0),
            self.memory(input_size),
        )

class API:
    def __init__(self):
        self.leaderboard = [] # {"filename": "name.pth", "name": "NAME", "surname": "surname"}*
        self.processing_list = {}
        self.min_max = {}
        self.model_processing_queue = Queue()
        self.lock = Lock()
        self.motd = {'type': 'info', 'message': 'Challenge 02 motd will appear here'}
        
        log.success("Starting process thread")
        processing_thread = Thread(target=self.compute_models)
        processing_thread.daemon = True
        processing_thread.start()

    def update_min_max(self):
        pp  = [e.perplexity    for e in self.leaderboard]
        mem = [e.memory        for e in self.leaderboard]
        per = [e.performance   for e in self.leaderboard]

        self.min_max['perplexity']      = (min(pp),  max(pp)  + 1)
        self.min_max['memory']          = (min(mem), max(mem) + 1)
        self.min_max['performance']     = (min(per), max(per) + 1)

    def load_leaderboard(self):
        pass
    
    def save_leaderboard(self):
        pass

    def get_leaderboard(self):
        with self.lock:
            return self.leaderboard

    def get_processing_list(self):
        return self.processing_list

    def run(self):
        log.success("Starting server")
        self.app.run(host="0.0.0.0", port=8080)

    def reset_leaderboard(self):
        with self.lock:
            self.leaderboard = []

    # Thread Target
    def compute_models(self):
        log.fine('Preparing dataset')
        # dataset = CIFAR10(
        #     root="./dataset", train=False, transform=T.ToTensor(), download=True,
        # )
        # loader = DataLoader(
        #     cifar10,
        #     batch_size=BATCH_SIZE,
        #     shuffle=False,
        #     num_workers=NUM_WORKERS,
        #     pin_memory=True,
        # )

        with open(VOCAB, "rb") as f:
            hpb = HarryPotterBooks(*pickle.load(f))
        hpd = HarryPotterDataset(hpb, window=MAX_SEQ)
        loader = DataLoader(
            hpd,
            batch_size=BATCH_SIZE,
            shuffle=False,
            num_workers=NUM_WORKERS,
            pin_memory=True,
        )

        log.success('Ready to process models')
        while True:
            filename = self.model_processing_queue.get()
            log.fine(f'Processing {filename}')
            self.processing_list[filename] = 'Processing'
            try:
                entry  = Entry(filename)
                record = entry(hpd, loader, MAX_SEQ, RUNS)
            except:
                log.error(f'Cannot process {filename}')
                traceback.print_exc()
                self.processing_list[filename] = 'Error, check your model'
                continue
            self.update_ranking(record, filename)

    def submit_file(self, filename: str):
        with self.lock:
            self.model_processing_queue.put(filename)
            self.processing_list[filename] = 'Waiting'

    def update_ranking(self, new_elem: Record, filename: str = None):
        with self.lock:
            if filename in self.processing_list:
                del self.processing_list[filename]

            # Delete previous record
            for e in self.leaderboard:
                if e.name == new_elem.name and e.surname == new_elem.surname:
                    self.leaderboard.remove(e)
                    break
            
            self.leaderboard.append(new_elem)
            self.update_min_max()
            for e in self.leaderboard:
                e.update_score_unscaled(self.min_max)
            σ = np.std([e.score_unscaled for e in self.leaderboard])
            for e in self.leaderboard:
                e.scale_score(σ)
            self.leaderboard = sorted(self.leaderboard, key=lambda d: -d.score)
            log.success("Ranking updated")


app = Flask(__name__)
CORS(app)
api = API()
allowed_pattern = re.compile("Challenge02_([A-z0-9]*)_([A-z0-9]*)([.]ts)")

@app.route('/leaderboard')
def get_leaderboard():
    return jsonify(api.get_leaderboard())

@app.route('/processing')
def get_processing():
    return api.get_processing_list()

def allowed_file(filename):
    m = allowed_pattern.match(filename)
    log.fine(f'Checking {filename}')
    return not (m == None or m.span()[1] != len(filename))

@app.route('/reset')
def reset_leaderboard():
    api.reset_leaderboard()
    return {}

@app.route('/motd', methods=['GET', 'POST'])
def manage_motd():
    if request.method == 'POST':
        if 'motd' in request.form:
            api.motd['message'] = request.form['motd']
        if 'type' in request.form:
            api.motd['type'] = request.form['type']
        return {}
    else:
        return api.motd

@app.route('/submit', methods=['POST'])
def submit():
    """
    Submit a file for evaluation
    """

    # Upload file
    if 'file' not in request.files:
        return {'error': 'No file'}, 400
    file = request.files['file']
    if not file or file.filename == '':
        log.error(f'No file')
        return {'error': 'Filename empty'}, 400
    if not allowed_file(file.filename):
        log.error(f'{file.filename} rejected')
        api.processing_list[file.filename] = 'Filename rejected'
        return {'error': 'Filename not allowed'}, 400

    filename = secure_filename(file.filename)
    path = os.path.join(FILE_UPLOAD_FOLDER, filename) # Complete relative filepath
    file.save(path)

    # Submit
    api.submit_file(path)
    return {}, 200

if __name__ == "__main__":
    os.makedirs(FILE_UPLOAD_FOLDER, exist_ok=True)
    app.run(host = "0.0.0.0", port = 8080)