"""frag_model."""
import numpy as np
import copy
import torch
import pytorch_lightning as pl
import torch.nn as nn
import torch_scatter as ts
import dgl.nn as dgl_nn

import ms_pred.common as common
import ms_pred.dag_pred.dag_data as dag_data
import ms_pred.nn_utils as nn_utils
import ms_pred.magma.fragmentation as fragmentation
import ms_pred.magma.run_magma as run_magma

class IntenGNN(pl.LightningModule):
    def __init__(
        self,
        hidden_size: int,
        gnn_layers: int = 2,
        mlp_layers: int = 0,
        set_layers: int = 2,
        learning_rate: float = 7e-4,
        lr_decay_rate: float = 1.0,
        weight_decay: float = 0,
        dropout: float = 0,
        mpnn_type: str = "PNA",
        pool_op: str = "avg",
        node_feats: int = common.ELEMENT_DIM + common.MAX_H,
        pe_embed_k: int = 0,
        max_broken: int = run_magma.FRAGMENT_ENGINE_PARAMS["max_broken_bonds"],
        frag_set_layers: int = 0,
        loss_fn: str = "cosine",
        root_encode: str = "gnn",
        inject_early: bool = False,
        warmup: int = 1000,
        embed_adduct=False,
        binned_targs: bool = True,
        encode_forms: bool = False,
        add_hs: bool = False,
        **kwargs,
    ):
        super().__init__()
        self.save_hyperparameters()
        self.hidden_size = hidden_size
        self.pe_embed_k = pe_embed_k
        self.root_encode = root_encode
        self.pool_op = pool_op
        self.inject_early = inject_early
        self.embed_adduct = embed_adduct
        self.binned_targs = binned_targs
        self.encode_forms = encode_forms
        self.add_hs = add_hs

        self.tree_processor = dag_data.TreeProcessor(
            root_encode=root_encode, pe_embed_k=pe_embed_k, add_hs=add_hs
        )

        self.formula_in_dim = 0
        if self.encode_forms:
            self.embedder = nn_utils.get_embedder("abs-sines")
            self.formula_dim = common.NORM_VEC.shape[0]

            # Calculate formula dim
            self.formula_in_dim = self.formula_dim * self.embedder.num_dim

            # Account for diffs
            self.formula_in_dim *= 2

        self.gnn_layers = gnn_layers
        self.set_layers = set_layers
        self.frag_set_layers = frag_set_layers
        self.mpnn_type = mpnn_type
        self.mlp_layers = mlp_layers

        self.learning_rate = learning_rate
        self.lr_decay_rate = lr_decay_rate
        self.warmup = warmup

        self.weight_decay = weight_decay
        self.dropout = dropout

        self.max_broken = max_broken + 1
        self.broken_onehot = torch.nn.Parameter(torch.eye(self.max_broken))
        self.broken_onehot.requires_grad = False
        self.broken_clamp = max_broken

        edge_feats = fragmentation.MAX_BONDS

        orig_node_feats = node_feats
        if self.inject_early:
            node_feats = node_feats + self.hidden_size

        adduct_shift = 0
        if self.embed_adduct:
            adduct_types = len(common.ion2onehot_pos)
            onehot = torch.eye(adduct_types)
            self.adduct_embedder = nn.Parameter(onehot.float())
            self.adduct_embedder.requires_grad = False
            adduct_shift = adduct_types

        # Define network
        self.gnn = nn_utils.MoleculeGNN(
            hidden_size=self.hidden_size,
            num_step_message_passing=self.gnn_layers,
            set_transform_layers=self.set_layers,
            mpnn_type=self.mpnn_type,
            gnn_node_feats=node_feats + adduct_shift,
            gnn_edge_feats=edge_feats,
            dropout=self.dropout,
        )

        if self.root_encode == "gnn":
            self.root_module = self.gnn

            # if inject early, need separate root and child GNN's
            if self.inject_early:
                self.root_module = nn_utils.MoleculeGNN(
                    hidden_size=self.hidden_size,
                    num_step_message_passing=self.gnn_layers,
                    set_transform_layers=self.set_layers,
                    mpnn_type=self.mpnn_type,
                    gnn_node_feats=node_feats + adduct_shift,
                    gnn_edge_feats=edge_feats,
                    dropout=self.dropout,
                )
        elif self.root_encode == "fp":
            self.root_module = nn_utils.MLPBlocks(
                input_size=2048,
                hidden_size=self.hidden_size,
                output_size=None,
                dropout=self.dropout,
                num_layers=1,
                use_residuals=True,
            )
        else:
            raise ValueError()

        # MLP layer to take representations from the pooling layer
        # And predict a single scalar value at each of them
        # I.e., Go from size B x 2h -> B x 1
        self.intermediate_out = nn_utils.MLPBlocks(
            input_size=self.hidden_size * 3 + self.max_broken + self.formula_in_dim,
            hidden_size=self.hidden_size,
            output_size=self.hidden_size,
            dropout=self.dropout,
            num_layers=self.mlp_layers,
            use_residuals=True,
        )

        trans_layer = nn_utils.TransformerEncoderLayer(
            self.hidden_size,
            nhead=8,
            batch_first=True,
            norm_first=False,
            dim_feedforward=self.hidden_size * 4,
        )
        self.trans_layers = nn_utils.get_clones(trans_layer, self.frag_set_layers)

        self.loss_fn_name = loss_fn
        if loss_fn == "mse":
            raise NotImplementedError()
        elif loss_fn == "cosine":
            self.loss_fn = self.cos_loss
            self.cos_fn = nn.CosineSimilarity()
            self.output_activations = [nn.Sigmoid()]
        else:
            raise NotImplementedError()

        self.num_outputs = len(self.output_activations)
        self.output_size = run_magma.FRAGMENT_ENGINE_PARAMS["max_broken_bonds"] * 2 + 1
        self.output_map = nn.Linear(
            self.hidden_size, self.num_outputs * self.output_size
        )

        # Define map from output layer to attn
        self.isomer_attn_out = copy.deepcopy(self.output_map)

        # Define buckets
        buckets = torch.DoubleTensor(np.linspace(0, 1500, 15000))
        self.inten_buckets = nn.Parameter(buckets)
        self.inten_buckets.requires_grad = False

        if self.pool_op == "avg":
            self.pool = dgl_nn.AvgPooling()
        elif self.pool_op == "attn":
            self.pool = dgl_nn.GlobalAttentionPooling(nn.Linear(hidden_size, 1))
        else:
            raise NotImplementedError()

        self.sigmoid = nn.Sigmoid()

    def cos_loss(self, pred, targ):
        """cos_loss.

        Args:
            pred:
            targ:
            optionlens:
        """
        if not self.binned_targs:
            raise ValueError()
        loss = 1 - self.cos_fn(pred, targ)
        loss = loss.mean()
        return {"loss": loss}

    def predict(
        self,
        graphs,
        root_reprs,
        ind_maps,
        num_frags,
        max_breaks,
        adducts,
        max_add_hs=None,
        max_remove_hs=None,
        masses=None,
        root_forms=None,
        frag_forms=None,
        binned_out=False,
    ) -> dict:
        """predict.

        Args:
            graphs:
            root_reprs:
            ind_maps:
            num_frags:
            max_breaks:
            adducts
            max_add_hs:
            max_remove_hs:
            masses
            root_forms
            frag_forms
            binned_out

        Returns:
            dict:
        """
        # B x nodes x num outputs x inten items
        out = self.forward(
            graphs,
            root_reprs,
            ind_maps,
            num_frags,
            adducts=adducts,
            broken=max_breaks,
            max_add_hs=max_add_hs,
            max_remove_hs=max_remove_hs,
            masses=masses,
            root_forms=root_forms,
            frag_forms=frag_forms,
        )

        if self.loss_fn_name not in ["mse", "cosine"]:
            raise NotImplementedError()

        output = out["output"][
            :,
            :,
            0,
        ]
        output_binned = out["output_binned"][:, 0, :]
        out_preds_binned = [i.cpu().detach().numpy() for i in output_binned]
        out_preds = [
            pred[:num_frag, :].cpu().detach().numpy()
            for pred, num_frag in zip(output, num_frags)
        ]

        if binned_out:
            out_dict = {
                "spec": out_preds_binned,
            }
        else:
            out_dict = {
                "spec": out_preds,
            }
        return out_dict

    def forward(
        self,
        graphs,
        root_repr,
        ind_maps,
        num_frags,
        broken,
        adducts,
        max_add_hs=None,
        max_remove_hs=None,
        masses=None,
        root_forms=None,
        frag_forms=None,
    ):
        """predict spec from graphs

        Args:
            graphs: DGL Graphs of all mols in batch
            root_reprs
            ind_maps: Mapping to pair graphs to root graphs (allows parallel)
            num_frags: Number of fragments
            broken: num bonds broken
            adducts
            max_add_hs
            max_remove_hs
            masses
            root_forms
            frag_forms
        """
        device = num_frags.device

        # if root fingerprints:
        embed_adducts = self.adduct_embedder[adducts.long()]
        if self.root_encode == "fp":
            root_embeddings = self.root_module(root_repr)
            raise NotImplementedError()
        elif self.root_encode == "gnn":
            with root_repr.local_scope():
                if self.embed_adduct:
                    embed_adducts_expand = embed_adducts.repeat_interleave(
                        root_repr.batch_num_nodes(), 0
                    )
                    ndata = root_repr.ndata["h"]
                    ndata = torch.cat([ndata, embed_adducts_expand], -1)
                    root_repr.ndata["h"] = ndata
                root_embeddings = self.root_module(root_repr)
                root_embeddings = self.pool(root_repr, root_embeddings)
        else:
            pass

        # Line up the features to be parallel between fragment avgs and root
        # graphs
        ext_root = root_embeddings[ind_maps]
        # Extend the root further to cover each individual atom
        ext_root_atoms = torch.repeat_interleave(
            ext_root, graphs.batch_num_nodes(), dim=0
        )
        concat_list = [graphs.ndata["h"]]

        if self.inject_early:
            concat_list.append(ext_root_atoms)

        if self.embed_adduct:
            adducts_mapped = embed_adducts[ind_maps]
            adducts_exp = torch.repeat_interleave(
                adducts_mapped, graphs.batch_num_nodes(), dim=0
            )
            concat_list.append(adducts_exp)

        with graphs.local_scope():
            graphs.ndata["h"] = torch.cat(concat_list, -1).float()

            frag_embeddings = self.gnn(graphs)

            # Average embed the full root molecules and fragments
            avg_frags = self.pool(graphs, frag_embeddings)

        # expand broken and map it to each fragment
        broken_arange = torch.arange(broken.shape[-1]).to(device)
        broken_mask = broken_arange[None, :] < num_frags[:, None]

        broken = torch.clamp(broken[broken_mask], max=self.broken_clamp)
        broken_onehots = self.broken_onehot[broken.long()]

        ### Build hidden with forms
        mlp_cat_list = [ext_root, ext_root - avg_frags, avg_frags, broken_onehots]

        hidden = torch.cat(mlp_cat_list, dim=1)

        # Pack so we can use interpeak attn
        padded_hidden = nn_utils.pad_packed_tensor(hidden, num_frags, 0)

        if self.encode_forms:
            diffs = root_forms[:, None, :] - frag_forms
            form_encodings = self.embedder(frag_forms)
            diff_encodings = self.embedder(diffs)
            new_hidden = torch.cat(
                [padded_hidden, form_encodings, diff_encodings], dim=-1
            )
            padded_hidden = new_hidden

        padded_hidden = self.intermediate_out(padded_hidden)
        batch_size, max_frags, hidden_dim = padded_hidden.shape

        # Build up a mask
        arange_frags = torch.arange(padded_hidden.shape[1]).to(device)
        attn_mask = ~(arange_frags[None, :] < num_frags[:, None])

        hidden = padded_hidden
        for trans_layer in self.trans_layers:
            hidden, _ = trans_layer(hidden, src_key_padding_mask=attn_mask)

        # hidden: B x L x h
        # attn_mask: B x L

        # Build mask
        max_inten_shift = (self.output_size - 1) / 2
        max_break_ar = torch.arange(self.output_size, device=device)[None, None, :].to(
            device
        )
        max_breaks_ub = max_add_hs + max_inten_shift
        max_breaks_lb = -max_remove_hs + max_inten_shift

        ub_mask = max_break_ar <= max_breaks_ub[:, :, None]
        lb_mask = max_break_ar >= max_breaks_lb[:, :, None]

        # B x Length x Mass shifts
        valid_pos = torch.logical_and(ub_mask, lb_mask)

        # B x Length x outputs x Mass shift
        valid_pos = valid_pos[:, :, None, :].expand(
            batch_size, max_frags, self.num_outputs, self.output_size
        )
        # B x L x Output
        output = self.output_map(hidden)
        attn_weights = self.isomer_attn_out(hidden)

        # B x L x Out x Mass shifts
        output = output.reshape(batch_size, max_frags, self.num_outputs, -1)
        attn_weights = attn_weights.reshape(batch_size, max_frags, self.num_outputs, -1)

        # Mask attn weights
        # attn_weights.masked_fill_(~valid_pos, -float("inf"))
        attn_weights.masked_fill_(~valid_pos, -99999)  # -float("inf"))

        # B x Out x L x Mass shifts
        output = output.transpose(1, 2)
        attn_weights = attn_weights.transpose(1, 2)
        valid_pos_binned = valid_pos.transpose(1, 2)

        # Calc inverse indices => B x Out x L x shift
        inverse_indices = torch.bucketize(masses, self.inten_buckets, right=False)
        inverse_indices = inverse_indices[:, None, :, :].expand(attn_weights.shape)

        # B x Out x (L * Mass shifts)
        attn_weights = attn_weights.reshape(batch_size, self.num_outputs, -1)
        output = output.reshape(batch_size, self.num_outputs, -1)
        inverse_indices = inverse_indices.reshape(batch_size, self.num_outputs, -1)
        valid_pos_binned = valid_pos.reshape(batch_size, self.num_outputs, -1)

        # B x Outs x ( L * mass shifts )
        pool_weights = ts.scatter_softmax(attn_weights, index=inverse_indices, dim=-1)
        weighted_out = pool_weights * output

        # B x Outs x (UNIQUE(L * mass shifts))
        output_binned = ts.scatter_add(
            weighted_out,
            index=inverse_indices,
            dim=-1,
            dim_size=self.inten_buckets.shape[-1],
        )

        output = output.reshape(batch_size, max_frags, self.num_outputs, -1)
        pool_weights_reshaped = pool_weights.reshape(
            batch_size, max_frags, self.num_outputs, -1
        )
        inverse_indices_reshaped = inverse_indices.reshape(
            batch_size, max_frags, self.num_outputs, -1
        )

        # B x Outs x binned
        valid_pos_binned = ts.scatter_max(
            (valid_pos_binned).long(),
            index=inverse_indices,
            dim_size=self.inten_buckets.shape[-1],
            dim=-1,
        )[0].bool()

        # Activate each dim with its respective output activation
        # Helpful for hurdle or probabilistic models
        new_outputs_binned = []
        for output_ind, act in enumerate(self.output_activations):
            new_outputs_binned.append(
                act(output_binned[:, output_ind : output_ind + 1, :])
            )
        output_binned = torch.cat(new_outputs_binned, -2)
        output_binned.masked_fill_(~valid_pos_binned, 0)

        # Index into output binned using inverse_indices_reshaped
        # Revert the binned output back to frags for attribution
        # B x Out x L x Mass shifts
        inverse_indices_reshaped_temp = inverse_indices_reshaped.transpose(
            1, 2
        ).reshape(batch_size, self.num_outputs, -1)
        output_unbinned = torch.take_along_dim(
            output_binned, inverse_indices_reshaped_temp, dim=-1
        )
        output_unbinned = output_unbinned.reshape(
            batch_size, self.num_outputs, max_frags, -1
        ).transpose(1, 2)
        output_unbinned_alpha = output_unbinned * pool_weights_reshaped

        return {"output_binned": output_binned, "output": output_unbinned_alpha}

    def _common_step(self, batch, name="train"):
        pred_obj = self.forward(
            batch["frag_graphs"],
            batch["root_reprs"],
            batch["inds"],
            batch["num_frags"],
            broken=batch["broken_bonds"],
            adducts=batch["adducts"],
            max_remove_hs=batch["max_remove_hs"],
            max_add_hs=batch["max_add_hs"],
            masses=batch["masses"],
            root_forms=batch["root_form_vecs"],
            frag_forms=batch["frag_form_vecs"],
        )
        pred_inten = pred_obj["output_binned"]
        pred_inten = pred_inten[:, 0, :]
        batch_size = len(batch["names"])

        loss = self.loss_fn(pred_inten, batch["inten_targs"])
        self.log(
            f"{name}_loss", loss["loss"].item(), batch_size=batch_size, on_epoch=True
        )

        for k, v in loss.items():
            if k != "loss":
                self.log(f"{name}_aux_{k}", v.item(), batch_size=batch_size)
        return loss

    def training_step(self, batch, batch_idx):
        """training_step."""
        return self._common_step(batch, name="train")

    def validation_step(self, batch, batch_idx):
        """validation_step."""
        return self._common_step(batch, name="val")

    def test_step(self, batch, batch_idx):
        """test_step."""
        return self._common_step(batch, name="test")

    def configure_optimizers(self):
        """configure_optimizers."""
        optimizer = torch.optim.Adam(
            self.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay
        )
        scheduler = nn_utils.build_lr_scheduler(
            optimizer=optimizer, lr_decay_rate=self.lr_decay_rate, warmup=self.warmup
        )
        ret = {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "frequency": 1,
                "interval": "step",
            },
        }
        return ret


from typing import Optional, Sequence

import numpy as np
import multiprocessing as mp

from npne import random

__all__ = ["SimpleEvaluator"]

class SimpleEvaluator:
    """
    Simple on-premise multiprocessing evaluator. Implement `evaluate_solution` to
    implement a new evaluator that inherits this class.

    Parameters:
        num_workers (Optional[int]): Number of parallel workers; if not specified, use
            the max number of available CPU cores (default: None).
    """

    def __init__(self, num_workers: Optional[int] = None) -> None:
        if num_workers is None:
            num_workers = mp.cpu_count()
        self.num_workers = num_workers

    def evaluate_solution(self, seed: int, solution: np.ndarray) -> None:
        """
        Evaluate a single solution candidate and return its fitness score. Implement
        a new evaluator class by implementing this method.

        Parameters:
            seed (int): A random seed for ensuring reproducibility of evaluation.
            solution (np.ndarray): A solution candidate represented with a NumPy array.

        Raises:
            NotImplementedError
        """
        raise NotImplementedError

    def evaluate(self, solutions: Sequence[np.ndarray]) -> np.ndarray:
        """
        Evaluate the argument set of solutions in subprocesses and return their fitness
        scores in a vector. Note that this method will fail if `evaluate_solution()` is
        not implemented.

        Parameters:
            solutions (Sequence[np.ndarray]): A sequence of solution candidates.

        Returns:
            A NumPy array of fitness scores.
        """
        results = []
        seeds = random.randint(2 ** 31 - 1, size=len(solutions)).tolist()
        with mp.Pool(self.num_workers) as pool:
            for seed, solution in zip(seeds, solutions):
                result = pool.apply_async(self.evaluate_solution, args=(seed, solution))
                results.append(result)
            fitness = [r.get() for r in results]
        return np.array(fitness)


import os

#import numpy as np

#from tkinter import E
import numpy as np
from argparse import ArgumentParser
import os
import torchvision

import torch
from models import load_model, train, evaluate #, evaluate
from datasets import AICityDatasetDetector, collate_dicts_fn

WAIT_TIME = 1
SAVE = False
CAR_LABEL_NUM = 3
FRAME_25PER = 510
EXPERIMENTS_FOLDER = "experiments"

def finetune(architecture_name, dataset_path, sequences, run_name, use_gpu=True):
    """
    Finetune object detector
    """
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    print('device:', device)

    transformations = torchvision.transforms.ToTensor()
    # transformations = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])

    test_dataset = AICityDatasetDetector(dataset_path, sequences, transformations=transformations)
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1, num_workers=1, shuffle=False, collate_fn=collate_dicts_fn)
    print('loader created')

    model, device = load_model(architecture_name, use_gpu, finetune=True)
    print('model loaded')

    model_folder_files = os.path.join(EXPERIMENTS_FOLDER, run_name)

    if not os.path.exists(EXPERIMENTS_FOLDER):
        # os.mkdir(EXPERIMENTS_FOLDER)
        os.makedirs(EXPERIMENTS_FOLDER,exist_ok=True)
    if not os.path.exists(model_folder_files):
        # os.mkdir(model_folder_files)
        os.makedirs(model_folder_files,exist_ok=True)

    ckpt_path = os.path.join(model_folder_files, run_name+"_best.ckpt")
    if os.path.exists(ckpt_path):
        model.load_state_dict(torch.load(ckpt_path))
    else:
        print("No finetuned weights for this experiment name, using pretrained model...")
    model.eval()

    log_bool = True
    num_epochs = 3
    batch_size = 1

    print('Starting mAP evaluation')
    evaluate(model, test_loader, device)
    print("Evaluation done")

def parse_arguments():
    parser = ArgumentParser()
    parser.add_argument("-d",
                        dest="dataset_path",
                        required=True,
                        type=str,
                        help="Path of the dataset.")
    parser.add_argument("-s",
                        dest="sequencies",
                        required=True,
                        type=str,
                        nargs='+')
    parser.add_argument("-n",
                        dest="architecture_name",
                        required=True,
                        type=str,
                        help="Architecture name. Options: FasterRCNN / MaskRCNN / ... ")
    parser.add_argument("-r",
                        dest="run_name",
                        required=True,
                        type=str,
                        help="Name of the experiment")
    parser.add_argument("-g",
                        default=False,
                        dest="use_gpu",
                        action="store_true",
                        help="Use GPU for model inference")
    args = parser.parse_args()

    return args.dataset_path, args.sequencies, args.architecture_name, args.use_gpu, args.run_name

if __name__ == "__main__":
    path_dataset, sequencies, architecture_name, use_gpu, run_name= parse_arguments()
    # (architecture_name, video_path, annotations, run_name, finetune, train_model=False, use_gpu=True)
    finetune(architecture_name, path_dataset, sequencies, run_name, use_gpu=use_gpu)


# posture_bed_data.jsonをposture_bed_data.csvに変換
import json
import csv

# ファイルの読み込み
with open("posture_bed_data.json", "r") as f: # 姿勢別のデータの読み込み
    posture_bed_data = json.load(f)

with open('tester_name.txt') as f: # 被験者名の読み込み
    tester_names = f.read().splitlines()

with open('tag_name.txt', 'r') as f: # タグ名の読み込み
    tag_names = f.read().splitlines()

# 列名の設定
header = ["posture", "tester", "E280116060000204AC6AD0EC", "E280116060000204AC6AD0E6", "E280116060000204AC6AD1FE", "E280116060000204AC6AD1FD", "E280116060000204AC6AC8F0" ,"E280116060000204AC6AD1FC"]

# csvファイルへの書き込み
with open("posture_bed_data.csv", "w") as f:
    writer = csv.writer(f)
    writer.writerow(header)
    for posture, posture_data in posture_bed_data.items():
        for i, tester in enumerate(tester_names):
            for j in range(len(posture_data[tester]["E280116060000204AC6AD0EC"])):
                data = [posture, i]
                for tag in tag_names:
                    data.append(posture_data[tester][tag][j])
                writer.writerow(data)


"""
Logistic regression on Sam's featurization.
"""

import scipy.io
import sklearn.linear_model
import sklearn.metrics
import sklearn.model_selection

data = scipy.io.loadmat('../../data/all_molecules_sam_feats.mat')
x, y = data['X'], data['Y'].ravel()

# Should match the split we've been using
X_train, X_test, y_train, y_test = sklearn.model_selection.train_test_split(
    x, y, test_size=0.25, random_state=int(sum(map(ord, 'alvin'))))

scipy.io.savemat('../../data/train_molecules_sam_feats.mat', {'X': X_train, 'Y': y_train})
scipy.io.savemat('../../data/test_molecules_sam_feats.mat', {'X': X_test, 'Y': y_test})

lr = sklearn.linear_model.LogisticRegression(verbose=1)
lr.fit(X_train, y_train)

y_pred = lr.predict(X_test)
print(sklearn.metrics.accuracy_score(y_pred, y_test))
print(sklearn.metrics.confusion_matrix(y_pred, y_test))


import os
os.environ["CUDA_VISIBLE_DEVICES"] = ""

import nobuco
from nobuco import ChannelOrder, ChannelOrderingStrategy
from nobuco.layers.weight import WeightLayer

import tensorflow as tf
import tensorflow_addons
from tensorflow.lite.python.lite import TFLiteConverter
from tensorflow import keras

import torch
from torch import nn, Tensor
import torch.nn.functional as F

class DenseImageWarp(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, image, flow):
        b, _, h, w = flow.shape

        lin_h = torch.linspace(start=0., end=(h - 1.), steps=h)[:, None].repeat(1, w)
        lin_w = torch.linspace(start=0., end=(w - 1.), steps=w)[None, :].repeat(h, 1)
        lin = torch.stack([lin_h, lin_w], dim=0)[None, ...]

        scale = torch.tensor([(h-1)/2, (w-1)/2])[None, None, None, :]

        grid = lin - flow
        grid = grid.permute(0, 2, 3, 1)
        grid = grid / scale - 1
        grid = torch.roll(grid, shifts=1, dims=-1)
        return F.grid_sample(image, grid, mode='bilinear', padding_mode='border', align_corners=True)

@nobuco.converter(DenseImageWarp, channel_ordering_strategy=ChannelOrderingStrategy.FORCE_TENSORFLOW_ORDER)
def dense_image_warp(self, image: Tensor, flow: Tensor):
    def func(image, flow):
        return keras.layers.Lambda(lambda args: tensorflow_addons.image.dense_image_warp(args[0], args[1]))([image, flow])
    return func

h = 256
w = 256
image = torch.rand(size=(1, 3, h, w))
flow = torch.rand(size=(1, 2, h, w))

inputs = [image, flow]
pytorch_module = DenseImageWarp().eval()

keras_model = nobuco.pytorch_to_keras(
    pytorch_module, inputs,
    inputs_channel_order=ChannelOrder.TENSORFLOW,
)

model_path = 'dense_image_warp'
keras_model.save(model_path + '.h5')
print('Model saved')

custom_objects = {'WeightLayer': WeightLayer}

keras_model_restored = keras.models.load_model(model_path + '.h5', custom_objects=custom_objects)
print('Model loaded')

converter = TFLiteConverter.from_keras_model_file(model_path + '.h5', custom_objects=custom_objects)
converter.target_ops = [tf.lite.OpsSet.SELECT_TF_OPS, tf.lite.OpsSet.TFLITE_BUILTINS]
tflite_model = converter.convert()
with open(model_path + '.tflite', 'wb') as f:
    f.write(tflite_model)


# coding: utf-8

# In[1]:
"""
def loadLibFolder (folder):
    import os, sys
    if folder not in sys.path:
        sys.path.insert(1, os.path.join(sys.path[0], folder))

# # Experimenting with POS dependency parser
# To be able to predict a category out of a sentence/text it is assumed that the POS tags and the dependency tree could have an inpact on the result. Here we investigate that relation

# In[2]:

from urllib import request, parse
import json
url = 'http://localhost:1337/sentence/'

# ## Sample text to try out the parser

# In[3]:

def parseSentence(sentence):
    try:
        sentence = request.quote(sentence)
        f =  request.urlopen(url + sentence)
        res = json.loads(f.read().decode('latin1'))
        return res
    except:
        return {'sentenceData': []}
def onlyNounsAndVerbs(data):
    return {
        'sentenceData': [word for word in data['sentenceData'] if 'NN' in word['tag'].split('|') or 'VB' in word['tag'].split('|')]
    }
def untilLevel(level, data):
    return {
        'sentenceData': [word for word in data['sentenceData'] if (int)(word['parent']) <= level]
    }
def toWordArray(data):
    return [word['base_word'] for word in data['sentenceData']]

# In[4]:

res = parseSentence('Han ler mot henne och hela hans ansikte säger att han älskar henne med hela sitt hjärta')

# In[5]:

# Example filtering
print ("Raw data:")
print (res)
print ("All words:")
print ([word['word'] for word in res['sentenceData']])
print ("Level three data:")
print ([word['word']+ '::' + word['tag'].split('|')[0] for word in res['sentenceData'] if (int)(word['parent']) <= 3])
print ("Only nouns and verbs:")
print ([word['word'] for word in res['sentenceData'] if 'NN' in word['tag'].split('|') or 'VB' in word['tag'].split('|')])

print(" ".join(toWordArray(untilLevel(3, onlyNounsAndVerbs(parseSentence(res))))))

# ## Classification experiment

# In[7]:

loadLibFolder('../gensim')

import os
import gensim
import gensim_documents
import dotenv
import numpy as np
dotenv.load()

# In[ ]:

limit_per_category = 2000
use_cache = False
use_all_data = True

# In[ ]:

categories = []
x_data = []
y_data = []
model = gensim.models.Doc2Vec.load(dotenv.get('DOC2VEC_MODEL'))

if use_cache and os.path.isfile('data/tmp_dependency_data'):
    with open('data/tmp_dependency_data', 'r', encoding='utf-8', errors='ignore') as tmp_cache_file:
        for category in tmp_cache_file:
            category = category[:-1]
            if category == "\n": continue
            if category not in categories:
                print ("TT", category)
                categories.append(category)
            sentVecs = []
            while True:
                sentence = tmp_cache_file.readline()[:-1]
                if sentence == "":
                    break
                artvec = model.infer_vector(doc_words=sentence.split())
                sentVecs.append(gensim.matutils.unitvec(artvec))
            y_data.append(categories.index(category))
            x_data.append(sentVecs)
else:
    data = gensim_documents.MMDBDocumentLists(dotenv.get('ARTICLE_PATH', '.') + '/csv_by_category/', useHeading=True, limit=limit_per_category)
    # with open('data/tmp_dependency_data', 'w', encoding='utf-8', errors='ignore') as tmp_cache_file:
    for i, doc in enumerate(data):
        if not doc.category in categories:
            categories.append(doc.category)
        #tmp_cache_file.write(doc.category + "\n")

        sentences = doc.content.split(".")
        sentVecs = []
        for j in range(20):
            if j >= len(sentences):
                sentVecs.append(np.zeros(300))
                continue
            if use_all_data:
                sentence = sentences[j]
            else:
                sentence = " ".join(toWordArray(untilLevel(3, onlyNounsAndVerbs(parseSentence(sentences[j])))))
            if sentence == "":
                sentVecs.append(np.zeros(300))
                continue
            artvec = model.infer_vector(doc_words=sentence.split())
            sentVecs.append(gensim.matutils.unitvec(artvec))
            #tmp_cache_file.write(sentence + "\n")
        #tmp_cache_file.write("\n")
        x_data.append(sentVecs)
        y_data.append(categories.index(doc.category))

        if i % (limit_per_category/4) == 0:
            print ("New epoch started, nr.", i+1, " of ", len(categories) * limit_per_category, " epochs")

# In[ ]:

"""

# #### Encode one hot vectors for the classes

# In[72]:
x_data = []
y_data = []
from random import shuffle
articles_train_text = open('segment/data/lstm_train/doc2Vec/train.txt','r',encoding='utf-8').read().split('\n')
articles_train_label = open('segment/data/lstm_train/doc2Vec/train_label2.txt','r',encoding='utf-8').read().split('\n')
articles_train = list(zip(articles_train_text, articles_train_label))
shuffle(articles_train)
articles_test_text = open('segment/data/lstm_train/doc2Vec/test.txt','r',encoding='utf-8').read().split('\n')
articles_test_label = open('segment/data/lstm_train/doc2Vec/test_label2.txt','r',encoding='utf-8').read().split('\n')
articles_test = list(zip(articles_test_text, articles_test_label))
shuffle(articles_test)

wordVectorLength = 250
docVectorLength = 100
import gensim
import numpy as np
categories = set(articles_train_label)
# Load word2vec model
#w2v_model = gensim.models.Word2Vec.load('segment/category_7_3.model')
d2v_model = gensim.models.Doc2Vec.load('segment/doc2vec.model')
"""
articles_train_labels, articles_train_vectors = zip(*[
    (int(label), [w2v_model.wv[word]
      for word in text.split(' ') if word in w2v_model.wv])
      for text, label in articles_train
])
articles_train_labels, articles_test_vectors = zip(*[
    (int(label), [w2v_model.wv[word]
      for word in text.split(' ') if word in w2v_model.wv])
      for text, label in articles_test
])
"""
import re
for text , label in articles_train:
    #sentences = re.split(r'。\?？!！',text)
    sentences = text.split('。')
    sentVecs = []
    #print(len(sentences))
    for j in range(20):
        if j >= len(sentences):
            sentVecs.append(np.zeros(100))
            continue
        else:
            sentence =  re.sub('。？?！!',' ',sentences[j].strip())
        if sentence == "":
            sentVecs.append(np.zeros(100))
            continue
        artvec = d2v_model.infer_vector(doc_words=sentence.split())
        sentVecs.append(gensim.matutils.unitvec(artvec))
    x_data.append(sentVecs)
    y_data.append(int(label))

for text , label in articles_test:
    sentences = text.split('。')
    sentVecs = []
    for j in range(20):
        if j >= len(sentences):
            sentVecs.append(np.zeros(100))
            continue
        else:
            sentence =  re.sub('。？?！!',' ',sentences[j].strip())
        if sentence == "":
            sentVecs.append(np.zeros(100))
            continue
        artvec = d2v_model.infer_vector(doc_words=sentence.split())
        sentVecs.append(gensim.matutils.unitvec(artvec))
    x_data.append(sentVecs)
    y_data.append(int(label))

###

"""
print ('out: ')
print(len([word for article in articles_train_text for word in article.split(' ') if not word in w2v_model.wv]))
print ('in')
print(len([word for article in articles_train_text for word in article.split(' ') if  word in w2v_model.wv]))

articles_train_vectors = [[article[i] if len(article) > i else np.zeros(wordVectorLength) for i in range(30)] for article in articles_train_vectors]
articles_test_vectors = [[article[i] if len(article) > i else np.zeros(wordVectorLength) for i in range(30)] for article in articles_train_vectors]
"""

"""
y_data_one_hot = np.zeros((len(articles_train_vectors), len(categories)))
y_data_one_hot[np.arange(len(articles_train_labels)), np.array(articles_train_labels)] = 1

x_data = articles_train_vectors
articles_train_vectors[0]
"""
categories = set(articles_train_label)
y_data_one_hot = np.zeros((len(y_data), len(categories)))
y_data_one_hot[np.arange(len(y_data)), np.array(y_data)] = 1
# ### LSTM classification with keras LSTM cells

# In[73]:

import keras
from keras.models import Sequential
from keras.layers import Dense, Dropout, Flatten
from keras.layers import Embedding
from keras.layers import LSTM
import numpy as np

# In[74]:

data_dim = len(x_data[0][0])
timesteps = len(x_data[0])
num_classes = len(categories)

# In[75]:

split = 0.4
limit_train = (int)(len(x_data) * split)
# Generate dummy training data
x_train = x_data[:limit_train]
y_train = y_data_one_hot[:limit_train]

# Generate dummy validation data
x_val = x_data[limit_train:]
y_val = y_data_one_hot[limit_train:]

print('matrix = 0 :')
# In[82]:
print(len([word for article in x_train for word in article if sum(word) == 0]))
print('matrix != 0 :')
# In[82]:
print(len([word for article in x_train for word in article if sum(word) != 0]))

print(len(x_val))
print(len(x_val[0]))
print(len(x_val[0][0]))

print(len(x_train))
print(len(x_train[0]))
print(len(x_train[0][0]))
print(len(categories))

print(y_train)

# To train a Sequential LSTM model that can classify a stacked sequence of words we need to define the input as follows:
#  * batch_size - number of datapoints in the dataset
#  * timesteps - the number of words per sequence
#  * data_dim - the number of features per word instance

# In[83]:

# expected input data shape: (batch_size, timesteps, data_dim)
model = Sequential()
model.add(LSTM(50, return_sequences=True,
               input_shape=(timesteps, data_dim)))  # returns a sequence of vectors of dimension 32
#model.add(Dropout(0.2))
model.add(LSTM(50, return_sequences=True))  # returns a sequence of vectors of dimension 32
model.add(LSTM(50))  # return a single vector of dimension 32
#model.add(Dropout(0.2))
model.add(Dense(num_classes, activation='softmax'))
#model.add(Dense(10, activation='softmax'))
#model.add(Flatten())
model.compile(loss='categorical_crossentropy',
              optimizer='rmsprop',
              metrics=['accuracy'])

model.fit(x_train, y_train, epochs=5, validation_data=(x_val, y_val))

# In[78]:

model.evaluate(x_val, y_val)

# In[84]:

model.predict([x_val[1]])

# In[ ]:

y_val[1]

# In[ ]:


from django.db import models
from django.contrib.auth.models import User
from datetime import datetime
from climats.models import Timekeeper, Company, Activity, Client, Case, Task
from smart_selects.db_fields import ChainedForeignKey
from django.http import request
from django.core.urlresolvers import reverse
from smart_selects.db_fields import ChainedForeignKey
from tlucidity.get_userobj import get_userobj

STATUSES = (
    ('O', 'Open'),
    ('R', 'Released'),
    ('E', 'Exported'),
)

class Entry(models.Model):
    user = models.ForeignKey(User, null=True)
    who = models.ForeignKey('climats.Timekeeper', verbose_name="Timekeeper", null=True, blank=True)
    work_date = models.DateField(null=True)
    company = models.ForeignKey('climats.Company', blank=True)
    matter = models.CharField(max_length=10, null=True)
    client = ChainedForeignKey(
        'climats.Client',
        chained_field="company",
        chained_model_field="company",
        limit_choices_to={'status':'O'},
        show_all=False,
        null=True,
        blank=True
    )
    case = ChainedForeignKey(
        'climats.Case',
        chained_field="client",
        chained_model_field="client",
        limit_choices_to={'status':'O'},
        show_all=False,
        auto_choose=True,
        null=True,
        blank=True
    )
    activity_code1 = models.ForeignKey('climats.Activity', related_name = 'acode1', null=True, blank=True)
    activity_code2 = models.ForeignKey('climats.Activity', related_name = 'acode2', null=True, blank=True)
    hours = models.DecimalField(max_digits=9, decimal_places=2, default=0)
    narrative = models.TextField(blank=True)
    created_date = models.DateField(auto_now=True,null=True)
    last_change_date = models.DateField(auto_now=True, null=True, blank=True)
    released_date = models.DateField(null=True, blank=True)
    exported_date = models.DateField(null=True, blank=True)
    status = models.CharField(max_length=1, choices=STATUSES, default='O')
    released = models.BooleanField(default=False)
    exported = models.BooleanField(default=False)
    recent_files = models.ForeignKey('Matter_use', null=True, blank=True, help_text='Select a recent file to override company, client & case')
    matter_keyin = models.CharField(max_length=10, null=True, blank=True, verbose_name="Matter", help_text='Type a client-case number to override client and case')

    class Meta:
        verbose_name_plural ='entries'

    def __str__(self):
        return u'%s, %s, %s, %s, %s' % (self.who.code, self.work_date.strftime("%m/%d/%y"), self.company.code, self.matter, self.hours)

    def get_absolute_url(self):
        return reverse('entry-view', kwargs={'pk': self.id})

    def _get_matter(self):
        fullstring = self.case.code
        clp = fullstring[2:7]
        sstring = clp + '-' + fullstring[7:11]
        return sstring
    matter = property(_get_matter)

    def save(self, *args, **kwargs):
        usobj = get_userobj()
        self.user = usobj
        self.matter_keyin = ''
        exported_date = ''
        if not self.who:
            self.who = Profile.objects.get(user_id=usobj.id).for_whom
            self.recent_files = None
            if self.released == True:
                self.status = 'R'
                self.released_date = datetime.now()
            super(Entry, self).save(*args, **kwargs)
            Profile.objects.filter(user_id=usobj.id).update(mre_date=self.work_date)
            rm_id = self.who.code + self.case.code
            rm_who = self.who
            rm_case = self.case
            rm_days = Profile.objects.get(user_id=usobj.id).recent_days
            if Matter_use.objects.filter(worker_case=rm_id).exists():
                Matter_use.objects.filter(worker_case=rm_id).update(last_used=datetime.now())
            else:
                Matter_use.objects.create(worker_case=rm_id, timekeeper=rm_who, matter=rm_case, last_used=datetime.now())
            for used_case in Matter_use.objects.filter(timekeeper=rm_who):
                diff = datetime.now().date() - used_case.last_used
                if diff.days > rm_days:
                    Matter_use.delete(used_case)
        else:
#           ('who already set - updating existing')
            self.last_change_date = datetime.now()
            self.recent_files = None
            if self.released == True:
                if self.released_date == None:
                    self.status = 'R'
                    self.released_date = datetime.now()
            else:
                self.status = 'O'
                self.released_date = None
            super(Entry, self).save(*args, **kwargs)
            rm_id = self.who.code + self.case.code
            rm_who = self.who
            rm_case = self.case
            rm_days = Profile.objects.get(user_id=usobj.id).recent_days
            if Matter_use.objects.filter(worker_case=rm_id).exists():
                Matter_use.objects.filter(worker_case=rm_id).update(last_used=datetime.now())
            else:
                Matter_use.objects.create(worker_case=rm_id, timekeeper=rm_who, matter=rm_case, last_used=datetime.now())
        return reverse('dashboard')

class Profile(models.Model):
    user = models.OneToOneField(User, primary_key=True)
    for_whom = models.ForeignKey('climats.Timekeeper', limit_choices_to={'status':'A'}, verbose_name='entering for', null=True, blank=True)
    mre_date = models.DateField(null=True, blank=True)
    recent_days = models.IntegerField(default=15, verbose_name='Days', help_text='until files deleted from recent files list.')

    def __str__(self):
        return self.user.username

    def get_absolute_url(self):
        return reverse('set-w', kwargs={'pk': self.user.id})

class Matter_use(models.Model):
    worker_case = models.CharField(max_length=15, primary_key=True)
    timekeeper = models.ForeignKey('climats.Timekeeper', limit_choices_to={'status':'A'})
    matter = models.ForeignKey('climats.Case', limit_choices_to={'status':'O'})
    last_used = models.DateField()

    def __str__(self):
        return u'%s %s(%s) %s(%s) ' % (self.matter.company.code, self.matter.client.name, self.matter.client.number, self.matter.name, self.matter.number)

class Expthru(models.Model):
    user = models.OneToOneField(User, primary_key=True)
    thru_date = models.DateField(default=datetime.now(), verbose_name="Through Date")
    key_list = models.CommaSeparatedIntegerField(max_length=32768, null=True, blank=True)

    def save(self, *args, **kwargs):
        usobj = get_userobj()
        self.user = usobj
        super(Expthru, self).save(*args, **kwargs)
        return reverse('dashboard')

class Export_Error(models.Model):
    date_time = models.DateTimeField(auto_now=True)
    ran_by = models.CharField(max_length=39)
    results = models.TextField()


import sys
import pygame
from windowClass import *

#Constants
WIDTH, HEIGHT = 960, 540
BACKGROUND = (120, 120, 120)
FPS = 24 #Refresh rate

#Functions
def getEvents():
    global running
    for event in pygame.event.get(): #Returns list of events since last call
        if event.type == pygame.QUIT:
             running = False
        if event.type == pygame.MOUSEBUTTONDOWN:
            mousePos = pygame.mouse.get_pos()
            if mouseOnGrid(mousePos):
                clickCell(mousePos)
                #print(mousePos) #(Debugging)

def update():
    window.update()

def draw():
    screen.fill(BACKGROUND)
    window.draw()

#mOG & cC used to track cell click state
def mouseOnGrid(pos):
    if pos[0] > 80 and pos[0] < WIDTH - 80:
        if pos[1] > 120 and pos[1] < HEIGHT - 20:
            return True
        return False

def clickCell(pos):
    grid_pos = [pos[0] - 80, pos[1] - 120]
    grid_pos[0] = grid_pos[0] // 20 #X pos
    grid_pos[1] = grid_pos[1] // 20 #Y pos

    #Sets cell dead/alive when clicked
    if window.grid[grid_pos[1]][grid_pos[0]].alive:
        window.grid[grid_pos[1]][grid_pos[0]].alive = False
    else:
        window.grid[grid_pos[1]][grid_pos[0]].alive = True
    #print("Clicked") #(Debugging)

#Main
pygame.init()
screen = pygame.display.set_mode((WIDTH, HEIGHT))
clock = pygame.time.Clock()
window = Window(screen, 80, 120) #positioned center, bottom

running = True
#Checks these on a functions while "running"
while running:
    mousePos = pygame.mouse.get_pos()
    getEvents()
    update()
    draw()
    pygame.display.update()
    clock.tick(FPS)
pygame.quit()
sys.exit()


from collections import defaultdict
import numpy as np
import pandas as pd
import operator
import time
# from kmincutunionfind import karger_min_cut
from kmincut import contract, fast_min_cut, Graph

# class Graph:
#     def __init__(self):
#         self.vertices = set()
#         self.edges = []

# class Edge:
#     def __init__(self, src, dest):
#         self.src = src
#         self.dest = dest

def load_data():
    df = pd.read_csv('releasedates.csv', sep=',', header=None)
    datalist = np.array(df)
    return datalist

def get_operator(comparison_operator):
    comparison_operator_dict = {
        "==": operator.eq,
        ">": operator.gt,
        ">=": operator.ge,
        "<": operator.lt,
        "<=": operator.le}
    return comparison_operator_dict[comparison_operator]

def nested_loop_join(datalistR, datalistS, conditions, block_size, num_blocks, num_columns, g):
    output = []
    # each block
    for bR in range(0, num_blocks * block_size, block_size):
        for bS in range(0, num_blocks * block_size, block_size):
            # each tuple
            for tR in range(bR, bR+block_size):
                for tS in range(bS, bS+block_size):
                    sign = conditions[0][2]
                    left = conditions[0][0]
                    right = conditions[0][1]
                    # TODO: remove hard code
                    if datalistR[tR][1] == 2 and get_operator(sign)(datalistR[tR][left], datalistS[tS][right]):
                        tuple_r = datalistR[tR]
                        tuple_s = datalistS[tS]

                        # add egeds to graph
                        src = "r" + str(tuple_r[0]) + "," + str(tuple_r[1])
                        dest = "s" + str(tuple_s[0]) + "," + str(tuple_s[1])
                        # edge = Edge(src, dest)
                        # g.vertices.add(src)
                        # g.vertices.add(dest)
                        # g.edges.append(edge)
                        g[src].append(dest)
                        g[dest].append(src)

                        # update output
                        output.append(np.concatenate(
                            (tuple_r, tuple_s), axis=0))

    return np.array(output), len(g)

def join():
    start = time.time()
    datalist = load_data()
    print("running time load data", time.time()-start)
    datalist = datalist[0:200, :]
    num_rows = np.shape(datalist)[0]
    num_columns = np.shape(datalist)[1]

    # g = Graph()
    g = defaultdict(list)

    # No header
    block_size = 4
    num_blocks = (num_rows - 1) // block_size  # divide by ceiling

    # query
    # select *
    # from releasedate r1 join releasedate r2
    #     on r1.releasedate < r2.releasedate

    start = time.time()
    # TODO: make this flexible to accept any queries
    datalistR = datalist
    datalistS = datalist.copy()
    conditions = [[2, 2, "<"]]
    join_results, no_of_vertices = nested_loop_join(datalistR, datalistS, conditions, block_size,
                                                    num_blocks, num_columns, g)

    print("running time nested loop join", time.time()-start)
    graph_list = []
    for item in g.items():
        graphs = [item[0]] + item[1]
        graph_list.append(graphs)

    # print(graph_list)
    # print("no of vertices:", no_of_vertices)
    return join_results, join_results.shape, graph_list, no_of_vertices

start = time.time()
join_results, result_shape, g, no_of_vertices = join()
# print(join_results, result_shape)
k = 20
# print("\n\nCut found by Karger's randomized algo is {}".format(
#     karger_min_cut(g, k, no_of_vertices)))
# karger_min_cut(g, k, no_of_vertices)
end = time.time()
print("running time load + nested loop join", end-start)

start = time.time()
graph = Graph(g)
end = time.time()
print("running time construct graph", end-start)

m = graph.edge_count
n = graph.vertex_count
print("m:", m)
print("n:", n)

# print(graph.verts)

start = time.time()
# m, gout = fast_min_cut(graph, k)
gout, groups = contract(graph, k)
# print()
# print(gout.verts)
print(gout.edges)
print()
print(gout.parents)
print()
end = time.time()
print("running time min cut:", end-start, "\n")


# Supress unnecessary warnings so that presentation looks clean
import warnings
warnings.filterwarnings('ignore')

# import sys
# sys.path.append('transformer')
import torch
from torch import nn
import torch.nn.functional as F
from torchtext import data
from argparse import ArgumentParser
from torch.utils.tensorboard import SummaryWriter
import tqdm     # CLI progress bar
import pandas as pd
import numpy as np
import matplotlib.pylab as plt
from transformers import GPT2Tokenizer, GPT2LMHeadModel
from decoder import top_p_sampling, scheduled_sampler, top_k_sampling, greedy_sampling

batch_size = 1
use_cuda_if_available = True
board = SummaryWriter()

def main(arg):

    device = 'cuda' if use_cuda_if_available and torch.cuda.is_available() else 'cpu'
    print("Device: ", device)

    try:
        convtexts = pd.read_csv('toxic.tsv', sep='\t')
    except:
        convtexts = pd.read_csv('adversarial/toxic.tsv', sep='\t')
    convtexts = np.array(convtexts).tolist()
    train_data_loader = torch.utils.data.DataLoader(convtexts, batch_size=batch_size, shuffle=True)
    validation_data_loader = torch.utils.data.DataLoader(convtexts, batch_size=batch_size, shuffle=False)

    print('num of training examples: ', len(train_data_loader))
    print('num of test examples: ', len(validation_data_loader))
    if arg.max_length < 0:
        mx = max([max([len(si) for si in s]) for s,t in train_data_loader])
        mx = mx * 2
    else:
        mx = arg.max_length
    print('maximum seq length: ', mx)

    # Load pre-trained model tokenizer (vocabulary)
    tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
    model = GPT2LMHeadModel.from_pretrained('gpt2')

    # it = 0
    # for src, tar in train_data_loader:
    #     if it == 495:
    #         print(src, tar)
    #     it += 1

    if torch.cuda.device_count() > 1:
        print("Running on ", torch.cuda.device_count(), "GPU's")
        model = nn.DataParallel(model)
        model.to(device)
    else:
        model.to(device)

    optim = torch.optim.Adam(lr=arg.lr, params=model.parameters())
    sched = torch.optim.lr_scheduler.LambdaLR(optim, lambda i: min(i / (arg.lr_warmup / arg.batch_size), 1.0))
    criterion = nn.L1Loss()
    # Training loop
    seen = 0
    training_loss_list = []
    validation_loss_list = []
    end_of_sentence = ['.', '<|endoftext|>', '?', '!', '."', '?"', '!"']
    for epoch in range(arg.num_epochs):
        print("\n Epoch: ", epoch)
        model.train(True)
        training_loss = 0
        batch_count = 0
        for source_list, target_list in train_data_loader:
            batch_count += 1
            for index in range(len(source_list)):
                optim.zero_grad()
                source = source_list[index]
                target = target_list[index]
                print('__Source__', source)
                print('__Target__', target)
                source_tokens = tokenizer.encode(source)
                target_tokens = tokenizer.encode(target)
                indexed_tokens = source_tokens
                target_token = target_tokens[0]
                predicted_tokens = []
                loss = 0
                num_pred = 0
                while num_pred < 50:
                    num_pred += 1
                    tokens_tensor = torch.tensor([indexed_tokens])
                    tokens_tensor.to(device)
                    outputs = model(tokens_tensor)
                    predictions = outputs[0][0][-1]
                    pred_prob = F.softmax(predictions)
                    predicted_index = top_p_sampling(pred_prob)
                    predicted_word = tokenizer.decode([predicted_index])
                    if predicted_word in end_of_sentence:
                        break
                    indexed_tokens += [predicted_index]
                    predicted_tokens.append(predicted_index)

                    loss += criterion(pred_prob[predicted_index], pred_prob[target_token])

                print('Predicted =', tokenizer.decode(predicted_tokens))
                if target_token in predicted_tokens:
                    loss = 0

                training_loss += loss
                loss.backward()
                if arg.gradient_clipping > 0.0:
                    nn.utils.clip_grad_norm_(model.parameters(), arg.gradient_clipping)
                optim.step()

        print("epoch", epoch, "training_loss", training_loss)
        training_loss_list.append(training_loss)

        board.add_scalar('training loss', training_loss, epoch)

        """
        with torch.no_grad():
            model.eval()
            batch_count = 0
            validation_loss = 0
            for source_list, target_list in tqdm.tqdm(validation_data_loader):
                batch_count += 1
                for index in range(len(source_list)):
                    optim.zero_grad()
                    source = source_list[index]
                    target = target_list[index]
                    source_size = len(source.split())
                    if source_size > mx:
                        source = source[:, :mx]
                    tok_inp = tokenizer.encode(source)
                    tok_tar = tokenizer.encode(target)
                    for tar in tok_tar:
                        inp_tensor = torch.tensor([tok_inp])
                        tar = torch.tensor(tar)
                        inp_tensor = inp_tensor.to(device)
                        tar = tar.to(device)
                        output = model(inp_tensor)
                        pred = output[0][0][-1]
                        pred_prob = F.softmax(pred)
                        loss = F.nll_loss(torch.log(pred_prob).unsqueeze(0), tar.unsqueeze(0))
                        validation_loss += loss
                        pred_idx = top_p_sampling(pred_prob, p=arg.top_p)
                        tok_inp.append(pred_idx)
                        seen += inp_tensor.size(0)
                        batch_count += 1

                    print("Source Given: ", source)
                    predicted_text = tokenizer.decode(tok_inp)
                    print("Predicted: ", predicted_text)
                    print("Target: ", target)

            sched.step()
            print("epoch", epoch, "validation_loss", validation_loss)
            validation_loss_list.append(validation_loss)
            # Save loss in text file
            c = [training_loss_list, validation_loss_list]
            with open("loss.txt", "a") as file:
                for x in zip(*c):
                    file.write("{0}\t{1}\n".format(*x))
        """

    model_file = 'saved_model.pkl'
    torch.save(model, model_file)

    plt.figure(figsize = (10, 4))
    plt.subplot(1, 2, 1)
    plt.plot(validation_loss_list, 'bo-', label = 'val-loss')
    plt.plot(training_loss_list, 'ro-', label = 'train-loss')
    plt.grid('on')
    plt.ylabel('loss')
    plt.xlabel('epoch')
    plt.legend(['validation', 'training'], loc='upper right')
    plt.savefig("Loss_vs_epoch.png")

if __name__ == '__main__':
    parser = ArgumentParser()

    """
    Arguments:
        --num-epochs,   -e  Number of epochs
        --batch-size,   -b  Batch size
        --learn-rate,   -l  Learning rate
        --tb_dir,       -T  Tensorboard logging directory
        --final,        -f  Whether to run on the real test set (if not included, validation set is used)
        --max_pool,         Use max pooling in final classification layer
        --embedding     -E  Size of character embeddings
        --vocab-size    -V  Vocabulary size
        --max           -M  Max sequence length. Longer sequences are clipped (-1 for no limit)
        --heads         -H  Number of attention heads
        --depth         -d  Depth of the network (num of self-attention layers)
        --random-seed   -r  RNG seed
        --lr-warmup         Learning rate warmup
        --gradient-clipping Gradient clipping
        --dropout       -D  Dropout
    """

    parser.add_argument("-e", "--num-epochs",
                        dest="num_epochs",
                        help="Number of epochs.",
                        default=5, type=int)

    parser.add_argument("-b", "--batch-size",
                        dest="batch_size",
                        help="The batch size.",
                        default=4, type=int)

    parser.add_argument("-l", "--learn-rate",
                        dest="lr",
                        help="Learning rate",
                        default=0.001, type=float)

    parser.add_argument("-T", "--tb_dir", dest="tb_dir",
                        help="Tensorboard logging directory",
                        default='./runs')

    parser.add_argument("-f", "--final", dest="final",
                        help="Whether to run on the real test set (if not included, the validation set is used).",
                        action="store_true")

    parser.add_argument("--max-pool", dest="max_pool",
                        help="Use max pooling in the final classification layer.",
                        action="store_true")

    parser.add_argument("-E", "--embedding", dest="embedding_size",
                        help="Size of the character embeddings.",
                        default=128, type=int)

    parser.add_argument("-V", "--vocab-size", dest="vocab_size",
                        help="Number of words in the vocabulary.",
                        default=50_000, type=int)

    parser.add_argument("-M", "--max", dest="max_length",
                        help="Max sequence length. Longer sequences are clipped (-1 for no limit).",
                        default=512, type=int)

    parser.add_argument("-H", "--heads", dest="num_heads",
                        help="Number of attention heads.",
                        default=8, type=int)

    parser.add_argument("-d", "--depth", dest="depth",
                        help="Depth of the network (nr. of self-attention layers)",
                        default=6, type=int)

    parser.add_argument("-r", "--random-seed",
                        dest="seed",
                        help="RNG seed. Negative for random",
                        default=1, type=int)

    parser.add_argument("--lr-warmup",
                        dest="lr_warmup",
                        help="Learning rate warmup.",
                        default=10_000, type=int)

    parser.add_argument("--gradient-clipping",
                        dest="gradient_clipping",
                        help="Gradient clipping.",
                        default=1.0, type=float)

    parser.add_argument("-D", "--dropout",
                        dest="dropout",
                        help="Dropout rate.",
                        default=0.0, type=float)

    parser.add_argument("-p", "--top-p",
                        dest="top_p",
                        help="Top P sampling.",
                        default=0.8, type=float)
    parser.add_argument("-k", "--top-k",
                        dest="top_k",
                        help="Top K sampling.",
                        default=5, type=float)
    parser.add_argument("-bs", "--beam-size",
                        dest="beam_size",
                        help="Beam search size.",
                        default=10, type=float)

    options = parser.parse_args()
    print('OPTIONS', options)

    main(options)


# your_app/management/commands/import_library.py

import csv
from django.core.management.base import BaseCommand
from books.models import Book

class Command(BaseCommand):
    help = 'Imports data from CSV file to the Book model.'

    def add_arguments(self, parser):
        parser.add_argument('csv_file', type=str, help='Path to the CSV file to import.')

    def handle(self, *args, **options):
        csv_file_path = options['csv_file']

        with open(csv_file_path, 'r') as csvfile:
            reader = csv.DictReader(csvfile)
            for row in reader:
                #
                if row['datePublished'] == '':
                    datePublished=0
                else:
                    datePublished=int(row['datePublished'])

                bid=row['bid']
                name=row['name']
                author=row['author']
                about=row['about']

                bookEdition=row['bookEdition']
                bookFormat=row['bookFormat']
                img=row['img']

                Book.objects.create(
                    bid=bid,
                    title=name,
                    author=author,
                    about=about,
                    date_published=datePublished,
                    book_edition=bookEdition,
                    book_format=bookFormat,
                    img=img
                )

        self.stdout.write(self.style.SUCCESS('Data imported successfully.'))


# -*- coding: utf-8 -*-
# __author__ = 'qingchen.zhang'

import os,re
import shutil
import gevent
from Bio import SeqIO
import numpy as np
import pandas as pd
from biocluster.module import Module
from biocluster.core.exceptions import OptionError
from biocluster.config import Config
from mbio.packages.metaasv.common_function import link_dir,link_file

class DeblurDenoiseSplitModule(Module):
    """
    按照样本进行拆分，Deblur 降噪, 基于训练好的数据进行降噪
    """
    def __init__(self, work_id):
        super(DeblurDenoiseSplitModule, self).__init__(work_id)
        options = [
            {"name": "fastq_dir", "type": "infile", "format": "sequence.fastq_dir"},##fastq_dir文件夹
            {'name': 'ref_fasta', 'type': 'infile', 'format': 'sequence.fasta'},  # 参考fasta序列,自定义模式上传的fasta文件
            {'name': 'database', 'type': 'string'},  # 数据库选择  用于选择参考数据库
            {"name": "min_sample_number", "type": "int", "default": 0},## 过滤长度的值
            {"name": "coverage", "type": "int", "default": 90},  # 根据90%的水平进行筛选， 理论上选择最小的文件
            {"name": "truc_len", "type": "int", "default": 0},##过滤长度的值，默认不过滤，如果不过滤deblur必须要对齐序列
            {"name": "min_size", "type": "int", "default": 1},
            {"name": "min_reads", "type": "int", "default": 10},
        ]
        self.add_option(options)
        self.fastq_dir = os.path.join(self.work_dir, "fastq_dir")
        self.trim_data = self.add_tool("metaasv.select_trimlength")
        self.file_list = []
        self.run_tools = []
        self.merge_tools = []

    def check_options(self):
        """
        参数二次检查
        :return:
        """
        if not self.option("fastq_dir").is_set:
            raise OptionError("请输入fastq文件夹！")
        if not self.option("database"):
            raise OptionError("请输入数据库类型！")

    def select_coverage(self):
        """
        根据每个样本的全部数据和占全部数据百分比的长度得到每个样本的trim长度
        :return:
        """
        options = {
            "fastq_dir":self.option("fastq_dir"),
            "coverage": self.option("coverage")
            }
        self.trim_data.set_options(options)
        self.trim_data.on("end", self.run_deblur)
        self.trim_data.run()

    def run_deblur(self):
        """
        运行软件deblur，非qiime2流程
        :return:
        """
        self.logger.info("min_sample_number: {}".format(self.option("truc_len")))
        if self.option("truc_len") == 0:
            result_path = os.path.join(self.trim_data.output_dir, "all_result.xls")
            min_number = self.get_min_seq_length(result_path)
        else:
            min_number = self.option("truc_len")
        if os.path.exists(self.fastq_dir):
            shutil.rmtree(self.fastq_dir)
        os.mkdir(self.fastq_dir)
        input_dir = self.option("fastq_dir").prop["path"]
        listdirs = os.listdir(input_dir)
        for file in listdirs:
            self.deblur = self.add_tool("metaasv.deblur_denoise_single")
            file_path = os.path.join(input_dir, file)
            new_file_path = os.path.join(self.fastq_dir, file)
            link_file(file_path, new_file_path)
            options = {
                "input_fastq": new_file_path,
                "database": self.option("database"),
                # "trim_length": self.option("trim_length"),
            }
            if self.option("ref_fasta").is_set:
                self.data_se = os.path.join(self.convert_file.output_dir, "data_se.qza")
                options["input_qza"] = self.data_se
            if self.option("ref_fasta").is_set:
                options["ref_fasta"] = os.path.join(self.taxon_file.output_dir, "ref_fasta.qza")## 用参考的序列
            if self.option("truc_len") != 0:
                options["truc_len"] = self.option("truc_len")
            elif self.option("coverage") != 0:
                options["truc_len"] = min_number
            else:
                options["truc_len"] = 90
            if self.option("min_size"):
                options["min_size"] = self.option("min_size")
            if self.option("min_reads"):
                options["min_reads"] = self.option("min_reads")
            self.deblur.set_options(options)
            self.run_tools.append(self.deblur)
        if len(self.run_tools) > 1:
            self.on_rely(self.run_tools, self.run_merge_file)
        else:
            self.run_tools[0].on('end', self.run_merge_file)
        for tool in self.run_tools:
            tool.run()
            gevent.sleep(0)

    def run_merge_file(self):
        """
        分别对feature-table、fasta和stat统计文件进行合并,fa去重
        :return:
        """
        asv_table = os.path.join(self.work_dir, "asv_table")
        if os.path.exists(asv_table):
            shutil.rmtree(asv_table)
        os.mkdir(asv_table)
        asv_fasta = os.path.join(self.work_dir, "asv_fasta")
        if os.path.exists(asv_fasta):
            shutil.rmtree(asv_fasta)
        os.mkdir(asv_fasta)
        n = 1
        j = 1
        for tool in self.run_tools:
            list_dirs = sorted(os.listdir(tool.output_dir))
            for file in list_dirs:
                file_path = os.path.join(tool.output_dir, file)
                if re.search(r'ASV_reps\.fa', file):
                    new_file = os.path.join(asv_fasta, file + str(n))
                    link_file(file_path, new_file)
                    n += 1
                elif re.search(r'ASV_table\.xls', file):
                    new_file = os.path.join(asv_table, file + str(j))
                    link_file(file_path, new_file)
                    j += 1

        self.merge1 = self.add_tool("metaasv.merge_file")
        options = {
            "input_dir": asv_table,
            "type": "table"
            }
        self.merge1.set_options(options)
        self.merge_tools.append(self.merge1)
        self.merge2 = self.add_tool("metaasv.merge_file")
        options = {
            "input_dir": asv_fasta,
            "type": "fasta"
            }
        self.merge2.set_options(options)
        self.merge_tools.append(self.merge2)

        if len(self.merge_tools) > 1:
            self.on_rely(self.merge_tools, self.create_mongo_data)
        else:
            self.merge_tools[0].on('end', self.create_mongo_data)
        for tool in self.merge_tools:
            tool.run()
            gevent.sleep(1)

    def create_mongo_data(self):
        """
        将ASV_table.xls和降噪后的序列文件统计信息进行合并
        :return:
        """
        result_dada = os.path.join(self.work_dir, "result_dada")
        if os.path.exists(result_dada):
            shutil.rmtree(result_dada)
        os.mkdir(result_dada)
        for tool in self.merge_tools:
            for file in sorted(os.listdir(tool.output_dir)):
                file_path = os.path.join(tool.output_dir, file)
                new_file = os.path.join(result_dada, file)
                link_file(file_path, new_file)
        asv_fasta = os.path.join(result_dada, "ASV_reps.fasta")
        asv_table_origin = os.path.join(result_dada, "ASV_table.xls")
        asv_table = pd.read_table(asv_table_origin, sep="\t", header=0)

        new_gene = {}
        number = 1
        ### 这里先对asv表的丰度进行过滤，然后根据过滤后的asv再提取序列，替换名称
        with open(asv_table_origin, 'r') as f, open(os.path.join(self.output_dir, "ASV_table.xls"), 'w') as w, open(os.path.join(self.output_dir, "ASV_md5.xls"), 'w') as outf:
            outf.write("ASV ID\tmd5\n")
            for line in f:
                if line.startswith("#OTU ID"):
                    line = line.strip().split("\t")
                    line[0] = "ASV ID"
                    new_line = "\t".join(line)
                    w.write(new_line+"\n")
                else:
                    line = line.strip().split("\t")
                    asv_id = line[0]
                    gene_name = "ASV" + str(number)
                    if asv_id not in new_gene:
                        new_gene[asv_id] = gene_name
                    line[0] = gene_name
                    asv_list = [float(x) for x in line[1:]]
                    total_asv_size = int(float(np.sum(asv_list)))
                    # choose_list = filter(lambda x: x != 0.0, asv_list)
                    # if len(choose_list) >= 2: ###在这里筛选asv_size，每条asv在至少两个样本中存在才保留，否则过滤掉
                    # if 10 >= self.option("min_reads"):
                    #     min_reads = 10
                    # else:
                    #     min_reads = self.option("min_reads")
                    if total_asv_size >= int(self.option("min_reads")):## 在这里筛选asv_abundance，所有样本的总和大于等于10的asv才会保留，否则过滤
                        number += 1
                        w.write("\t".join(line) + "\n")
                        outf.write("{}\t{}\n".format(gene_name, asv_id))

        with open(os.path.join(self.output_dir, "ASV_reps.fasta"), "w") as wf:
            for seq_record in SeqIO.parse(asv_fasta, 'fasta'):
                gene_id = seq_record.id
                gene_seq = seq_record.seq
                if gene_id in new_gene:
                    gene_name = new_gene[gene_id]
                    wf.write(">{}\n{}\n".format(gene_name, gene_seq))

        mongo_denoise_stat = self.work_dir + "/denoise_stat.xls"
        sample_list = list(asv_table.columns)[1:]
        sample_number_dict = {}
        sequence_number_dict = {}
        for sample in sample_list:
            asv_number = len(list(asv_table[sample][asv_table[sample] != 0.0]))
            if sample not in sample_number_dict:
                sample_number_dict[sample] = asv_number
            sequence_number = asv_table[sample][asv_table[sample] != 0.0].values.sum()
            if sample not in sequence_number_dict:
                sequence_number_dict[sample] = sequence_number

        denoise_path = result_dada + "/Deblur_sequence_info.txt"
        with open(mongo_denoise_stat, 'w') as w:
            w.write("Sample_Name\tAsv_Number\tSeq_Number\n")
            for sample in sample_list:
                if sample in sample_number_dict:
                    asv_num = sample_number_dict[sample]
                if sample in sequence_number_dict:
                    valied_number = sequence_number_dict[sample]
                w.write("{}\t{}\t{}\n".format(sample, asv_num, valied_number))
        if os.path.exists(denoise_path):
            os.remove(denoise_path)
        if os.path.exists(mongo_denoise_stat):
            os.rename(mongo_denoise_stat, denoise_path)
        self.run_convert_qza()

    def run_convert_qza(self):
        """
        转table转为qza文件
        :return:
        """
        qza_table = os.path.join(self.work_dir, "result_dada")
        allfiles = sorted(os.listdir(qza_table))
        for file in allfiles:
            file_path = os.path.join(qza_table, file)
            self.convert_qza = self.add_tool("metaasv.file_to_qzaqzv")
            if re.search(r'ASV_table', file):
                options = {
                "input_file": file_path,
                "type": "table",
                "prefix": "ASV_table"
                }
            elif re.search(r'ASV_reps', file):
                options = {
                "input_file": file_path,
                "type": "fasta",
                "prefix": "ASV_reps"
                }
            else:
                options = {
                    "input_file": file_path,
                    "type": "table",
                    "prefix": "Deblur_stats"
                }
            self.convert_qza.set_options(options)
            self.file_list.append(self.convert_qza)
        if len(self.file_list) > 1:
            self.on_rely(self.file_list, self.set_output)
        else:
            self.file_list[0].on("end", self.set_output)
        for tool in self.file_list:
            tool.run()
            gevent.sleep(0)

    def get_min_seq_length(self, inputfilr):
        """
        根据传入的inputfile获得最小值
        :param inputfilr:
        :return:
        """
        min_number_list = []
        with open(inputfilr, "r") as f:
            for line in f:
                line = line.strip().split("\t")
                sample_min = int(float(line[1]))
                if sample_min not in min_number_list:
                    min_number_list.append(sample_min)
        min_number = min(min_number_list)
        return min_number

    def set_output(self):
        """
        设置结果文件目录
        :return:
        """
        self.logger.info("开始设置结果文件目录！")
        dada2_stat = os.path.join(self.output_dir, "Deblur_sequence_info.txt")
        if not os.path.exists(dada2_stat):
            link_file(os.path.join(self.work_dir,"result_dada", "Deblur_sequence_info.txt"),dada2_stat)
        for tool in self.file_list:
            for file in os.listdir(tool.output_dir):
                file_path = os.path.join(tool.output_dir,file)
                out_path = os.path.join(self.output_dir, file)
                link_file(file_path, out_path)
        # for module in self.merge_tools:
        #     for file in os.listdir(module.output_dir):
        #         file_path = os.path.join(module.output_dir,file)
        #         out_path = os.path.join(self.output_dir, file)
        #         link_file(file_path, out_path)
        self.logger.info("设置结果文件目录成功！")
        self.end()

    def run(self):
        """
        运行
        :return:
        """
        super(DeblurDenoiseSplitModule, self).run()
        if self.option("truc_len") == 0:
            self.select_coverage()
        else:
            self.run_deblur()

    def end(self):
        """
        结束
        :return:
        """
        super(DeblurDenoiseSplitModule, self).end()


# OOP Concepts

# What is class?
''' Class is like a template or blueprint. We use it to build objects. Class it itself does'nt occupyany memory . when an object
ia initiaited memmory is allocated '''

# What does a class contain?
''' A contain methods and attributes. Methods are equvelent to functions, Attributes are like variables  '''
'''All the variables you create in python are object.
When you wrie name='apple', you cerate an objrct(variable) belonging to the class string.
All the functions can do on strings like rstrip, lstrip, split etc are the methods in the class string'''

from email.headerregistry import Address

name='apple'
print(type(name))
li=[]
print(type(li))

# let us create a list, tuple, dictionary
l=[] #l is an object belonging to class list
t=()# t is an object belonging to classtuple
d={} # d is an object belonging to class dictionary
print(type(l))
print(type(t))
print(type(d))

# Let us see what all the methods or function that you can see in a given class. We have created 4 variables  till now

# name which is string. I which is a list, t which is a tuple , d which is a list.
print(dir(name))#to see the methods associated with class strings

namel =" This is Skill-Lync "
print(namel)
namel.strip( ' ')
na=namel.isdigit()
print(na)

print(dir(l))# to see methods associated with thw class list.
#kindly note the pressence of apped wich allow you to edit the list
# strings don't have it

print(dir(t))# to see methods associated with thw class tuple
print(dir(d))# to see methods associated with thw class dictionary

## now let us add soemthing to the list.
#we call the object and the method you want to perform

l=[]
print(l)
print(50*'*')
l.append(1)
l.append(2)
print(l)# you can see that class list has a method called apped
# Which allowed you to change the list.

print(name)

print('*'*50)
# name.append(name) #AttributeError: 'str' object has no attribute 'append'
print(name)## you can see that class string has a method called apped
# Which allowed you to change the string.
''' Note the error "str " object has no attribute "append" '''

#Let us create a class now

class Student:
    pass
stu1=Student()
stu2=Student()
stu3=Student()
stu4=Student()

print(stu1)
print(stu2)
print(stu3)
print(stu4)

# Let us complicate our Student class

class Student1:
    pass
stu1=Student()
stu2=Student()
stu3=Student()
stu4=Student()
q=Student1()
# creating  an instance variable
stu1.name='Abhi'
stu1.rollno=100
stu1.dob=1990
stu2.name='Ridhi'
stu2.rollno=181
stu2.dob=1999
stu3.name='kumar'
stu3.rollno=672
stu3.dob=2000
stu4.name='sandhya'
stu4.rollno=871
stu4.dob=2001
print('*'*100)
print(stu1.name,stu1.rollno,stu1.dob)
print(stu2.name,stu2.rollno,stu2.dob)
print(stu3.name,stu3.rollno,stu3.dob)
print(stu4.name,stu4.rollno,stu4.dob)

# This is extrememly painful if there are 100 Students. So use init Method.
# This is also called  a constructor in kanguages like JAVA

class Students:
    def __init__(self,name,rollno,dob,city):
        self.name=name
        self.rollno =rollno
        self.dob=dob
        self.city=city

stu1=Students('Abhishek',100,1990,'Bareilly')
stu2=Students('Anand',101,1999,'Lucknow')
stu3=Students('Rohit',102,2000,'Dehli')
stu4=Students('Mohit',104,2010,'Chennai')

print(stu1.name,stu1.rollno,stu1.dob,stu1.city)
print(stu2.name,stu2.rollno,stu2.dob,stu2.city)
print(stu3.name,stu3.rollno,stu3.dob,stu3.city)
print(stu4.name,stu4.rollno,stu4.dob,stu4.city)

# Creating Another Method

class Students1:
    def __init__(self,name,rollno,dob,city):
        self.name=name
        self.rollno =rollno
        self.dob=dob
        self.city=city
    def address(self):
        addres=f' \n Name : {self.name} ,\n Roll no: {self.rollno},\n DOb : {self.dob}, \n City : {self.city} '
        return addres

stu1=Students1('Abhishek',100,1990,'Bareilly')
stu2=Students1('Anand',101,1999,'Lucknow')
stu3=Students1('Rohit',102,2000,'Dehli')
stu4=Students1('Mohit',104,2010,'Chennai')

print(stu1.name,stu1.rollno,stu1.dob,stu1.city)
print(stu2.name,stu2.rollno,stu2.dob,stu2.city)
print(stu3.name,stu3.rollno,stu3.dob,stu3.city)
print(stu4.name,stu4.rollno,stu4.dob,stu4.city)

print('*'*50)
print(stu1.address())
print(stu2.address())
print(stu3.address())
print(stu4.address())

## Let us create another method to find age
#to do this we import the date module and get the current year from that

from datetime import date
class Students2:
    def __init__(self,name,rollno,dob,city):
        self.name=name
        self.rollno =rollno
        self.dob=dob
        self.city=city
    def address(self):
        addres=f' \n Name : {self.name} ,\n Roll no: {self.rollno},\n DOb : {self.dob}, \n City : {self.city} '
        return addres

    def age(self):
        return date.today().year-self.dob

stu1=Students2('Abhishek',100,1990,'Bareilly')
stu2=Students2('Anand',101,1999,'Lucknow')
stu3=Students2('Rohit',102,2000,'Dehli')
stu4=Students2('Mohit',104,2010,'Chennai')

print("\t AGE")
print('*'*50)
print(stu1.age())
print(stu2.age())
print(stu3.age())
print(stu4.age())
print('*'*50)

print('*'*50)
print(stu1.address())
print(stu2.rollno)

# Instance the variables, creating object specific attributes like say fees for student1 ect..

from datetime import date
class Students2:
    def __init__(self,name,rollno,dob,city):
        self.name=name
        self.rollno =rollno
        self.dob=dob
        self.city=city
    def address(self):
        addres=f' \n Name : {self.name} ,\n Roll no: {self.rollno},\n DOb : {self.dob}, \n City : {self.city} '
        return addres

    def age(self):
        return date.today().year-self.dob

stu1=Students2('Abhishek',100,1990,'Bareilly')
stu2=Students2('Anand',101,1999,'Lucknow')
stu3=Students2('Rohit',102,2000,'Dehli')
stu4=Students2('Mohit',104,2010,'Chennai')

stu1.fees=1300000
stu2.fees=1300000
stu3.fees=1300000
print(stu1.fees)
print(stu2.fees)
print(stu3.fees)
# print(stu4.fees) #This will throw an error , because no attributes for stu4
#AttributeError: 'Students2' object has no attribute 'fees'

## To create commmon fees for all studenrs we create the class variable outside  the class

from datetime import date
class Students3:
    def __init__(self,name,rollno,dob,city):
        self.name=name
        self.rollno =rollno
        self.dob=dob
        self.city=city
    def address(self):
        addres=f' \n Name : {self.name} ,\n Roll no: {self.rollno},\n DOb : {self.dob}, \n City : {self.city} '
        return addres

    def age(self):
        return date.today().year-self.dob

stu1=Students3('Abhishek',100,1990,'Bareilly')
stu2=Students3('Anand',101,1999,'Lucknow')
stu3=Students3('Rohit',102,2000,'Dehli')
stu4=Students3('Mohit',104,2010,'Chennai')

Students3.fees=150000 #  Comman Class Variables

print(stu1.fees)
print(stu2.fees)
print(stu3.fees)
print(stu4.fees)

#Now let us creae the class variable inside the class only

from datetime import date

class Students3:
    fees=10000 # class variables inside the class
    def __init__(self,name,rollno,dob,city):
        self.name=name
        self.rollno =rollno
        self.dob=dob
        self.city=city
    def address(self):
        addres=f' \n Name : {self.name} ,\n Roll no: {self.rollno},\n DOb : {self.dob}, \n City : {self.city} '
        return addres

    def age(self):
        return date.today().year-self.dob

stu1=Students3('Abhishek',100,1990,'Bareilly')
stu2=Students3('Anand',101,1999,'Lucknow')
stu3=Students3('Rohit',102,2000,'Dehli')
stu4=Students3('Mohit',104,2010,'Chennai')

print(stu1.fees)
print(stu2.fees)
print(stu3.fees)
print(stu4.fees)

# All the fees of students are 10000. now we can change one particular students fees alone by using the object directly.
print('*'*50)
print('# All the fees of students are 10000. now we can change one particular students fees alone by using the object directly.')
stu1.fees=120000
print(stu1.fees)
print(stu2.fees)
print(stu3.fees)
print(stu4.fees)


from collections import OrderedDict
import numpy as np
from math import ceil
import cv2
from io import BytesIO
from pathlib import Path
from PIL import Image

from utils import *
from huffman import *

def decode_header(reader: BytesIO):
    read = reader.read
    assert read(2) == MARKER.SOI  # Start of Image
    marker = read(2)
    if marker == MARKER.APP0:  # Application Marker 0
        APP0 = dict(
            length=read(2),
            identifier=read(5),
            version=read(2),  # b'\x01\x01'
            unit=read(1),
            x_density=read(2),
            y_density=read(2),
        )
        APP0['thumbnail-data'] = read(bytes2int(APP0['length']) - 14)
        marker = read(2)

    APPns = []
    while MARKER.APPn[0] <= marker <= MARKER.APPn[1]:  # Application Marker n
        APPn = dict(marker=marker, length=read(2))
        APPns.append(APPn)
        APPn['content'] = read(bytes2int(APPn['length']) - 2)
        marker = read(2)

    qts = {}
    while marker == MARKER.DQT:  # Define Quantization Table
        length, byte = bytes2int(read(2)), read(1)[0]
        precision, id_ = byte >> 4, byte & 0b1111
        assert precision == 0, 'only support uint8'
        table_size = (precision + 1) * 64
        assert length == table_size + 3
        dtype = np.uint8 if precision == 0 else np.uint16
        qts[id_] = np.frombuffer(
            read(table_size), dtype=dtype, count=64).astype(np.int32)
        marker = read(2)

    assert marker == MARKER.SOF0  # Start of Frame
    length = bytes2int(read(2))
    precision = read(1)[0]
    assert precision == 8, 'only support 8'
    height, width = bytes2int(read(2)), bytes2int(read(2))
    cop_num = read(1)[0]  # 1: grayscale; 3: YCrCb;
    assert cop_num in (1, 3), 'only support grayscale and RGB images'
    # components_name: ['Y', 'Cr', 'Cb']
    cop_infos = {}
    for i in range(cop_num):
        cop_id, ratio, qt_id = read(3)  # cop_id: 1, 2, 3; qt_id: 0, 1
        horizontal_ratio, vertical_ratio = ratio >> 4, ratio & 15
        cop_infos[cop_id] = ComponentInfo(
            id_=cop_id,
            horizontal=horizontal_ratio,
            vertical=vertical_ratio,
            qt_id=qt_id,
            dc_ht_id=None,
            ac_ht_id=None,
        )

    marker = read(2)
    hts = {}  # (type, id): table
    # type 0: DC, 1: AC
    # Huffman table id: 0, 1
    while marker == MARKER.DHT:  # Define Huffman Table
        length, byte = bytes2int(read(2)), read(1)[0]
        type_, id_ = byte >> 4, byte & 0b1111
        count = np.frombuffer(read(16), dtype=np.uint8, count=16)
        total = count.sum()
        weigh = np.frombuffer(read(total), dtype=np.uint8, count=total)
        assert length == 19 + total
        hts[type_, id_] = canonical_huffman_table(count, weigh)
        marker = read(2)

    DRI = None
    if marker == MARKER.DRI:  # Define Restart Interval
        DRI = dict(marker=marker, length=bytes2int(read(2)))
        DRI['content'] = read(length - 2)
        marker = read(2)
    assert DRI is None, 'not support DRI yet'

    assert marker == MARKER.SOS
    length = bytes2int(read(2))
    assert read(1)[0] == cop_num
    for i in range(cop_num):
        cop_id, ht_id = read(2)
        # ht_id: (dc_ht_id, ac_ht_id)
        cop_infos[cop_id].dc_ht_id = ht_id >> 4
        cop_infos[cop_id].ac_ht_id = ht_id & 0b1111
    # cop_id: 1 Y, 2 Cb, 3 Cr
    cop_infos = [cop_infos[i + 1] for i in range(cop_num)]
    assert read(3) == b'\x00\x3f\x00'
    return qts, hts, cop_infos, height, width

def decode_mcu(stream, mcu_num, block_num, hts):
    mcu = np.zeros([mcu_num, block_num, 64], dtype=np.int32)
    for cur in mcu:
        for dct, (dc_ht, ac_ht) in zip(cur, hts):
            # decode DC coefficients
            dc_size = read_huffman_code(dc_ht, stream)
            dc_complement = stream.read_int(dc_size)
            dct[0] = decode_2s_complement(dc_complement, dc_size)

            ac_index = 1
            rss_ = []
            while ac_index < 64:
                # obtain (run, size) pair
                rs = read_huffman_code(ac_ht, stream)
                run_length, ac_size = rs >> 4, rs & 15
                if (run_length, ac_size) == (0, 0):
                    dct[ac_index:] = 0
                    break
                dct[ac_index:ac_index + run_length] = 0
                ac_index += run_length
                ac_complement = stream.read_int(ac_size)
                dct[ac_index] = decode_2s_complement(ac_complement, ac_size)
                rss_.append((run_length, ac_size,
                             decode_2s_complement(ac_complement, ac_size)))
                ac_index += 1
    # Due to byte alignment, some bits (< 8) may be not used.
    return mcu

def iDPCM(mcu, cop_infos):
    index, mcu_num = 0, mcu.shape[0]
    for cop in cop_infos:
        num = cop.horizontal * cop.vertical
        dc = mcu[:, index:index + num, 0].flat
        for i in range(1, num * mcu_num):
            dc[i] += dc[i - 1]
        index += num
    return mcu

def de_quantization(mcu, cop_infos, qts):
    index = 0
    for cop in cop_infos:
        num = cop.horizontal * cop.vertical
        mcu[:, index:index + num] *= qts[cop.qt_id][None, None]
        index += num
    return mcu

def zigzag_decode(zz):
    dct = np.zeros((zz.shape[0], zz.shape[1], BH, BW), dtype=np.float32)
    trace = zigzag_points(BH, BW)
    for i, p in enumerate(trace):
        dct[:, :, p[0], p[1]] = zz[:, :, i]
    return dct

def iDCT(dct):
    block = np.zeros_like(dct)
    # the input of DCT/iDCT must be float number
    for i in range(dct.shape[0]):
        for j in range(dct.shape[1]):
            block[i, j] = cv2.idct(dct[i, j])
    return block

def decode_jpeg(reader):
    qts, hts, cop_infos, height, width = decode_header(reader)
    remains = reader.read()
    data, marker = remains[:-2].replace(b'\xff\x00', b'\xff'), remains[-2:]
    bit_stream = BitStreamReader(data)
    assert marker == MARKER.EOI

    # for blocks in MCU, choose the corresponding huffman tables to decode
    mcu_hts = []
    for cop in cop_infos:
        for _ in range(cop.horizontal * cop.vertical):
            mcu_hts.append((hts[0, cop.dc_ht_id], hts[1, cop.ac_ht_id]))

    # the height and width of MCU
    mh, mw = 8 * cop_infos[0].vertical, 8 * cop_infos[0].horizontal
    mcu_num = ceil(height / mh) * ceil(width / mw)
    block_num = len(mcu_hts)  # the number of blocks in each MCU
    mcu = decode_mcu(bit_stream, mcu_num, block_num, mcu_hts)

    # reverse the differential encoding (DPCM) of DC coefficients
    mcu = iDPCM(mcu, cop_infos)
    # inverse quantization
    mcu = de_quantization(mcu, cop_infos, qts)

    # zigzag to block, the returned result is float32
    mcu = zigzag_decode(mcu)

    # iDCT
    block = iDCT(mcu)

    # arrange blocks to form images
    ims, index = [], 0
    for cop in cop_infos:
        hr, vr = cop.horizontal, cop.vertical
        num = hr * vr
        # the number of MCUs in each column/row
        nh, nw = ceil(height / mh), ceil(width / mw)
        if num == 1:  # Cb, Cr, or Y when 4:4:4
            im = restore_image(block[:, index], nh, nw)
        else:  # Y when 4:2:2 or 4:2:0
            mat = block[:, index:index + num].reshape(
                -1, vr, hr, BH, BW).swapaxes(2, 3).reshape(-1, vr*BH, hr*BW)
            im = restore_image(mat, nh, nw)
        index += num
        ims.append(im)
    # reverse DC level shift for luminance
    ims[0] += 128
    if len(ims) == 3:
        # interpolating for chroma
        if ims[0].shape != ims[1].shape:
            h_, w_ = ims[0].shape
            for i in range(1, 3):
                ims[i] = cv2.resize(ims[i], (w_, h_))
        im = np.stack(ims, axis=-1)
        im = YCbCr2RGB(im)
    else:
        im = ims[0]
    im = im[:height, :width].round()
    im[im < 0], im[im > 255] = 0, 255
    return im.astype(np.uint8)

def read_jpeg(filename):
    reader = BytesIO(Path(filename).read_bytes())
    return decode_jpeg(reader)

def main():
    im = read_jpeg('image/color.jpg')
    Image.fromarray(im).save('image/color-decode.bmp')

if __name__ == '__main__':
    main()


import sys
import pygame
import basicGameLogic
from math import log2

# UI
size = width, height = 480, 500
playRegion = 480, 480
FPS = 60

# Colors
black = (0, 0, 0)
white = (255, 255, 255)
fontColor = (82, 52, 42)
defaultTileColor = (232, 232, 232)
tileBoarderColor = fontColor

# Game
boardSize = 4

def drawBoard(screen, board):
    screen.fill(black)
    for i in range(boardSize):
        for j in range(boardSize):
            color = defaultTileColor
            numberText = ''
            if board[i][j] != 0:
                gComponent = 235 - log2(board[i][j]) * ((235 - 52) / (boardSize ** 2))
                color = (235, gComponent, 52)
                numberText = str(board[i][j])
            rect = pygame.Rect(j * playRegion[0] / boardSize,
                                i * playRegion[1] / boardSize,
                                playRegion[0] / boardSize,
                                playRegion[1] / boardSize)

            pygame.draw.rect(screen, color, rect)
            pygame.draw.rect(screen, fontColor, rect, 1)

            fontImage = tileFont.render(numberText, 0, fontColor)
            if fontImage.get_width() > playRegion[0] / boardSize:
                fontImage = pygame.transform.scale(fontImage,
                                                   (playRegion[0] / boardSize,
                                                    fontImage.get_height() / fontImage.get_width() * playRegion[
                                                        0] / boardSize))
            screen.blit(fontImage,
                        (j * playRegion[0] / boardSize + (playRegion[0] / boardSize - fontImage.get_width()) / 2,
                         i * playRegion[1] / boardSize + (playRegion[1] / boardSize - fontImage.get_height()) / 2))

    fontImage = scoreFont.render("Score: {:,}".format(sum(map(sum, board))), 1, white)  # Sum of all elements
    screen.blit(fontImage, (1, playRegion[1] + 1))

def handleInput(event, board):
    if event.type == pygame.QUIT:
        sys.exit()
    if event.type == pygame.KEYDOWN:
        if event.key == pygame.K_RIGHT:
            basicGameLogic.board = basicGameLogic.merge_right(basicGameLogic.board)
            basicGameLogic.addNewValue()
        elif event.key == pygame.K_LEFT:
            basicGameLogic.board = basicGameLogic.merge_left(basicGameLogic.board)
            basicGameLogic.addNewValue()
        elif event.key == pygame.K_UP:
            basicGameLogic.board = basicGameLogic.merge_up(basicGameLogic.board)
            basicGameLogic.addNewValue()
        elif event.key == pygame.K_DOWN:
            basicGameLogic.board = basicGameLogic.merge_down(basicGameLogic.board)
            basicGameLogic.addNewValue()

def gameLoop():
    clock = pygame.time.Clock()
    basicGameLogic.board = basicGameLogic.initialize_board(boardSize)

    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                sys.exit()
            elif event.type == pygame.KEYDOWN:
                handleInput(event, basicGameLogic.board)

        drawBoard(screen, basicGameLogic.board)
        pygame.display.flip()
        clock.tick(FPS)

if __name__ == '__main__':
    pygame.init()
    screen = pygame.display.set_mode(size)
    pygame.display.set_caption("2048")
    tileFont = pygame.font.SysFont("", 72)
    scoreFont = pygame.font.SysFont("", 22)
    gameLoop()