# adapted from Distill-VQ
import os.path as osp
import numpy as np
from collections import defaultdict
from tqdm import tqdm
import faiss
import pickle
import torch
from torch import nn
import torch.nn.functional as F
import torch.distributed as dist
from scipy.optimize import linear_sum_assignment

class ProductQuantization(nn.Module):
    def __init__(
        self,
        pq_type: str = 'pq',
        subvector_num: int = 32,
        subvector_bits: int = 8,
        dist_mode: str = 'ip',
        emb_size: int = 768,
        pq_init_method: str = 'faiss',
        pq_update_method: str = 'grad',
        tie_nci_pq_centroid: int = 0,
        lm_head: nn.Parameter = None,
        centroid_update_loss: str = 'none',
        rq_topk_score: str = 'prod',
    ):
        super().__init__()
        assert pq_type in ('pq', 'opq', 'rq')
        assert dist_mode in ('ip', 'l2', 'iptol2')
        # dist_mode 'iptol2' means using l2 to represent ip
        # https://gist.github.com/mdouze/e4bdb404dbd976c83fe447e529e5c9dc
        # query (q1,...,qd) to (q1,...,qd,0)
        # document (p1,...,pd) to (p1,...,pd,sqrt(PHI-l2norm^2))
        self.pq_type = pq_type
        self.subvector_num = subvector_num
        self.subvector_bits = subvector_bits
        self.subvector_cents = 2 ** subvector_bits
        self.dist_mode = dist_mode
        self.emb_size = emb_size
        self.pq_init_method = pq_init_method
        self.pq_update_method = pq_update_method
        self.tie_nci_pq_centroid = tie_nci_pq_centroid
        self.centroid_update_loss = centroid_update_loss
        self.rq_topk_score = rq_topk_score
        self.get_preds = False

        if self.pq_type == 'rq':
            last_dim = emb_size
        else:
            last_dim = emb_size // subvector_num
        self.last_dim = last_dim

        if self.tie_nci_pq_centroid:
            assert self.dist_mode != 'iptol2'
            assert self.pq_type != 'opq'
            self.lm_head = lm_head
            if self.pq_type == 'pq':
                self.weight_layers = nn.Sequential(
                    nn.Linear(emb_size, emb_size), nn.ReLU(), nn.Linear(emb_size, last_dim))
        else:
            codebook_last_dim = last_dim
            if self.dist_mode == 'iptol2':
                codebook_last_dim += 1
            self.codebook = nn.Parameter(torch.empty(
                subvector_num, self.subvector_cents, codebook_last_dim), requires_grad=(self.pq_update_method == 'grad')).type(torch.FloatTensor)
            if self.pq_type == 'opq':
                self.rotate = nn.Parameter(torch.empty(
                    emb_size, emb_size), requires_grad=False).type(torch.FloatTensor)
            if self.pq_update_method == 'ema':
                # adapted from rq-vae-transformer
                self.decay = 0.99
                self.eps = 1e-5
                self.restart_unused_codes = True
                self.register_buffer('cluster_size_ema', torch.zeros(
                    subvector_num, self.subvector_cents))
                self.register_buffer(
                    'embed_ema', self.codebook.detach().clone())

    def augment_xb(self, xb, phi=None):
        norms = np.sum((xb ** 2), axis=-1)
        if phi is None:
            phi = np.max(norms)
        extracol = np.sqrt(phi - norms)
        return np.hstack((xb, extracol[..., np.newaxis]))

    def augment_xq(self, xq):
        if isinstance(xq, torch.Tensor):
            xq = torch.cat((xq, xq.new_zeros(*xq.shape[:-1], 1)), dim=-1)
        else:
            xq = np.concatenate(
                (xq, np.zeros((*xq.shape[:-1], 1), dtype=xq.dtype)), axis=-1)
        return xq

    def wrapped_augment_xb(self, xb, index=None):
        if self.dist_mode != 'iptol2':
            return xb
        emb_size = self.last_dim
        reshaped = False
        if xb.shape[-1] not in (emb_size, emb_size + 1):
            ori_shape = xb.shape
            xb = xb.reshape(-1, emb_size)
            reshaped = True
        if xb.shape[-1] == emb_size:
            xb = self.augment_xb(xb)
            extracol = xb[..., -1]
            if reshaped:
                xb = xb.reshape(*ori_shape[:-1], -1)
        else:
            extracol = xb[..., -1]
        assert xb.shape[-1] == emb_size + 1
        if index is None:
            self.extracol[:] = torch.tensor(extracol, dtype=torch.float32)
        else:
            self.extracol[:, index] = torch.tensor(
                extracol, dtype=torch.float32)
        return xb

    def rq_minus_centroids(self, embeddings, centroids):
        return embeddings - centroids[..., :self.last_dim]

    def compute_scores(self, a, b):
        if self.dist_mode == 'ip':
            result = a * b
        else:
            if self.dist_mode == 'iptol2' and a.shape[-1] != b.shape[-1]:
                a = self.augment_xq(a)
            result = - ((a - b) ** 2)
        return torch.sum(result, dim=-1)

    def get_codebook(self):
        if self.tie_nci_pq_centroid:
            codebook = self.lm_head[2:2+self.subvector_num*self.subvector_cents].reshape(
                self.subvector_num, self.subvector_cents, self.emb_size)
            if self.pq_type == 'pq':
                codebook = self.weight_layers(codebook)
            return codebook
        else:
            return self.codebook

    def codebook_from_index(self, index, index_file=None):
        if index is None:
            index = faiss.read_index(index_file)
        if self.pq_type == 'opq':
            assert isinstance(index, faiss.IndexPreTransform)
            vt = faiss.downcast_VectorTransform(index.chain.at(0))
            assert isinstance(vt, faiss.LinearTransform)
            rotate = faiss.vector_to_array(vt.A).reshape(vt.d_out, vt.d_in)
            index = faiss.downcast_index(index.index)
        else:
            rotate = None
            index = index

        if self.pq_type == 'rq':
            rq = index.rq
            centroid_embeds = faiss.vector_to_array(rq.codebooks)
            codebook = centroid_embeds.reshape(
                rq.M, rq.total_codebook_size // rq.M, -1)
            assert self.last_dim == rq.d
        else:
            pq_index = index
            centroid_embeds = faiss.vector_to_array(pq_index.pq.centroids)
            codebook = centroid_embeds.reshape(
                pq_index.pq.M, pq_index.pq.ksub, -1)
            assert self.last_dim == pq_index.pq.dsub
        assert self.codebook.shape[-1] == codebook.shape[-1], 'Codebook shape not aligned! Considering whether iptol2.'

        with torch.no_grad():
            self.codebook.copy_(torch.tensor(codebook))
            if self.pq_type == 'opq':
                self.rotate.copy_(torch.tensor(rotate))

    def build_faiss_index(self, doc_embeddings, save_path=None):
        assert self.dist_mode in ('ip', 'l2', 'iptol2')
        index_metric = faiss.METRIC_INNER_PRODUCT if self.dist_mode == 'ip' else faiss.METRIC_L2

        if self.pq_type == 'pq':
            components = f"PQ{self.subvector_num}x{self.subvector_bits}"
        elif self.pq_type == 'opq':
            components = f"OPQ{self.subvector_num},PQ{self.subvector_num}x{self.subvector_bits}"
        elif self.pq_type == 'rq':
            components = f"RQ{self.subvector_num}x{self.subvector_bits}"
        if self.dist_mode == 'iptol2':
            assert self.pq_type != 'rq', 'Not support iptol2 rq using faiss index.'
            doc_embeddings = self.wrapped_augment_xb(doc_embeddings)
        emb_size = doc_embeddings.shape[-1]

        index = faiss.index_factory(
            emb_size, components, index_metric)

        index.train(doc_embeddings)
        index.add(doc_embeddings)

        if save_path is not None:
            faiss.write_index(index, save_path)
        return index

    @torch.no_grad()
    def get_document_cluster_simple(self, return_mapping: bool = False):
        assert self.get_preds
        cluster = defaultdict(list)
        mapping = {}
        for i, p in enumerate(self.last_preds):
            label = tuple(p.tolist())
            mapping[i] = label
            cluster[label].append(i)
        del self.last_preds
        self.get_preds = False
        if return_mapping:
            return dict(cluster), mapping
        else:
            return dict(cluster)

    @torch.no_grad()
    def get_document_cluster(self, doc_embeddings: np.ndarray, rank: int, nrank: int, batch_size: int = 1024, return_mapping: bool = False):
        num_docs = doc_embeddings.shape[0]
        num_docs_per_worker = num_docs // nrank
        start = num_docs_per_worker * rank
        if rank + 1 == nrank:
            ending = num_docs
        else:
            ending = start + num_docs_per_worker
        part_num_docs = ending - start
        cluster = torch.empty(
            (part_num_docs, self.subvector_num), dtype=torch.int32)
        if self.pq_type == 'rq':
            func = self.get_rq_document_cluster
        else:
            func = self.get_pq_document_cluster
        func(doc_embeddings, cluster, start,
             ending, rank, batch_size)
        if return_mapping:
            new_mapping = {}
        doc_cluster = defaultdict(list)
        for k, v in enumerate(cluster):
            tuplev = tuple(v.tolist())
            key = k + start
            doc_cluster[tuplev].append(key)
            if return_mapping:
                new_mapping[key] = tuplev
        print('Number of document clusters:', len(doc_cluster))
        if return_mapping:
            return dict(doc_cluster), new_mapping
        else:
            return dict(doc_cluster)

    def get_pq_document_cluster(self, doc_embeddings: np.ndarray, cluster: torch.Tensor, start: int, ending: int, rank: int, batch_size: int = 1024):
        emb_size = doc_embeddings.shape[1]
        part_emb_size = emb_size // self.subvector_num
        codebook = self.get_codebook().cpu()
        with torch.no_grad():
            for i in tqdm(range(start, ending, batch_size), desc=f'Document Cluster {rank}'):
                batch_start = i
                batch_ending = min(i+batch_size, ending)
                cur_embeddings = torch.tensor(
                    doc_embeddings[batch_start:batch_ending], device='cpu')
                if self.pq_type == 'opq':
                    cur_embeddings = torch.matmul(
                        cur_embeddings, self.rotate.cpu().T)
                cur_codebook = codebook.unsqueeze(
                    1).expand(-1, cur_embeddings.size(0), -1, -1)
                if self.dist_mode == 'iptol2':
                    cur_extracol = self.extracol[batch_start:batch_ending]
                for j in range(self.subvector_num):
                    emb_start = j*part_emb_size
                    emb_ending = emb_start + part_emb_size
                    part_doc_embeddings = cur_embeddings[:,
                                                         emb_start:emb_ending]
                    if self.dist_mode == 'iptol2':
                        part_doc_embeddings = torch.cat(
                            (part_doc_embeddings, cur_extracol[:, j].unsqueeze(-1)), -1)
                    part_doc_embeddings = part_doc_embeddings.unsqueeze(-2)
                    part_codebook = cur_codebook[j]
                    distance = self.compute_scores(
                        part_doc_embeddings, part_codebook)
                    index = distance.max(dim=-1, keepdim=True)[1].squeeze(-1)
                    cluster[batch_start - start: batch_ending - start, j] = index

    def get_rq_document_cluster(self, doc_embeddings: np.ndarray, cluster: torch.Tensor, start: int, ending: int, rank: int, batch_size: int = 1024):
        part_num_docs = ending - start
        temp_docemb = torch.tensor(doc_embeddings[start:ending], device='cpu')
        codebook = self.get_codebook()
        with torch.no_grad():
            for j in range(self.subvector_num):
                part_codebook = codebook[j].cpu()
                if self.dist_mode == 'iptol2':
                    cur_extracol = self.extracol[:, j].unsqueeze(-1)
                for i in tqdm(range(0, part_num_docs, batch_size), desc=f'Document Cluster {rank}-{j}'):
                    batch_start = i
                    batch_ending = min(i+batch_size, part_num_docs)
                    cur_embeddings = temp_docemb[batch_start:batch_ending]
                    if self.dist_mode == 'iptol2':
                        cur_embeddings = torch.cat(
                            (cur_embeddings, cur_extracol[batch_start:batch_ending]), -1)
                    cur_embeddings = cur_embeddings.unsqueeze(-2)
                    cur_codebook = part_codebook.unsqueeze(
                        0).expand(cur_embeddings.size(0), -1, -1)
                    distance = self.compute_scores(
                        cur_embeddings, cur_codebook)
                    index = distance.max(dim=-1, keepdim=True)[1].squeeze(-1)
                    cluster[batch_start:batch_ending, j] = index
                    temp_docemb[batch_start:batch_ending] = self.rq_minus_centroids(
                        temp_docemb[batch_start:batch_ending], part_codebook[index])

    def forward(self, vecs, return_loss=True):
        # TODO: check results and determine which is better
        # before commit we generate label AFTER gumbel softmax
        # after commit we generate label BEFORE gumbel softmax
        if self.pq_type == 'rq':
            proba, index, loss = self.forward_rq(vecs, return_loss)
        else:
            proba, index, loss = self.forward_pq(vecs, return_loss)
        if self.training and self.pq_update_method == 'ema':
            self.ema_update(vecs, index)
        return proba, index, loss

    def forward_pq(self, vecs, return_loss=True):
        if self.pq_type == 'opq':
            vecs = torch.matmul(vecs, self.rotate.T)
        vecs = vecs.view(vecs.size(0), self.subvector_num, -1)
        codebook = self.get_codebook().unsqueeze(0).expand(vecs.size(0), -1, -1, -1)
        proba = self.compute_scores(vecs.unsqueeze(-2), codebook)
        index = proba.max(dim=-1)[1]
        if return_loss:
            if self.centroid_update_loss == 'reconstruct':
                reconstruct_emb = self.get_reconstruct_vector(index, codebook).view(
                    index.shape[0], self.subvector_num, self.last_dim)
                loss = ((vecs - reconstruct_emb) ** 2).mean()
            else:
                loss = None
        else:
            loss = None
        return proba, index, loss

    def forward_rq(self, vecs, return_loss=True):
        allproba = []
        index = []
        codebook = self.get_codebook()
        use_reconstruct_loss = self.centroid_update_loss == 'reconstruct'
        if use_reconstruct_loss:
            errors = []
        for i in range(self.subvector_num):
            cur_codebook = codebook[i:i+1].expand(vecs.size(0), -1, -1)
            proba = self.compute_scores(vecs.unsqueeze(-2), cur_codebook)
            part_index = proba.max(dim=-1)[1]
            allproba.append(proba)
            index.append(part_index)
            if self.dist_mode == 'iptol2':
                cur_centroid = codebook[i][part_index][..., :-1]
            else:
                cur_centroid = codebook[i][part_index]
            if use_reconstruct_loss:
                cur_error = vecs.detach() - cur_centroid
                errors.append(cur_error)
            if i != self.subvector_num - 1:
                vecs -= cur_centroid.detach()
        proba = torch.stack(allproba, dim=1)
        index = torch.stack(index, dim=1)
        if return_loss:
            if use_reconstruct_loss:
                loss = (torch.stack(errors) ** 2).mean()
                # loss = (vecs ** 2).mean()
            else:
                loss = None
        else:
            loss = None
        return proba, index, loss

    @torch.no_grad()
    def ema_update(self, vectors, idxs):
        assert self.dist_mode != 'iptol2'
        embed_dim = self.last_dim
        if self.pq_type == 'rq':
            vectors = vectors.unsqueeze(
                1).expand(-1, self.subvector_num, -1).reshape(-1, embed_dim)
        else:
            if self.pq_type == 'opq':
                vectors = torch.matmul(vectors, self.rotate.T)
            vectors = vectors.reshape(-1, embed_dim)
        idxs = idxs.reshape(-1)
        n_vectors = vectors.shape[0]
        one_hot_idxs = vectors.new_zeros(n_vectors, self.subvector_cents)
        one_hot_idxs.scatter_(dim=1,
                              index=idxs.unsqueeze(1),
                              src=vectors.new_ones(n_vectors, 1)
                              )
        one_hot_idxs = one_hot_idxs.reshape(-1,
                                            self.subvector_num, self.subvector_cents)
        cluster_size = one_hot_idxs.sum(dim=0)
        vectors_sum_per_cluster = torch.bmm(one_hot_idxs.permute(
            1, 2, 0), vectors.reshape(-1, self.subvector_num, embed_dim).permute(1, 0, 2))

        if dist.is_initialized():
            dist.all_reduce(vectors_sum_per_cluster, op=dist.ReduceOp.SUM)
            dist.all_reduce(cluster_size, op=dist.ReduceOp.SUM)

        self.cluster_size_ema.mul_(self.decay).add_(
            cluster_size, alpha=1 - self.decay)
        self.embed_ema.mul_(self.decay).add_(
            vectors_sum_per_cluster, alpha=1 - self.decay)

        if self.restart_unused_codes:
            temp = vectors.reshape(-1, self.subvector_num, embed_dim)
            B = temp.shape[0]
            if B < self.subvector_cents:
                n_repeats = (self.subvector_cents + B - 1) // B
                std = temp.new_ones(embed_dim) * 0.01 / np.sqrt(embed_dim)
                temp = temp.repeat(n_repeats, 1, 1)
                temp = temp + torch.rand_like(temp) * std

            _vectors_random = torch.stack([temp[torch.randperm(
                temp.shape[0], device=temp.device), i][:self.subvector_cents] for i in range(self.subvector_num)])

            if dist.is_initialized():
                dist.broadcast(_vectors_random, 0)

            usage = (self.cluster_size_ema.unsqueeze(-1) >= 1).float()
            self.embed_ema.mul_(usage).add_(_vectors_random * (1-usage))
            usage = usage.squeeze(-1)
            self.cluster_size_ema.mul_(usage)
            self.cluster_size_ema.add_(torch.ones_like(
                self.cluster_size_ema) * (1-usage))

        # update embedding
        n = self.cluster_size_ema.sum(dim=1, keepdim=True)
        normalized_cluster_size = (
            n * (self.cluster_size_ema + self.eps) /
            (n + self.subvector_cents * self.eps)
        )
        self.codebook.data[:] = self.embed_ema / \
            normalized_cluster_size.unsqueeze(-1)

    def fix(self):
        self.codebook.requires_grad_(False)
        if self.pq_type == 'opq':
            self.rotate.requires_grad_(False)

    @torch.no_grad()
    def initialize(self, index_file, doc_emb, rank, seed, pq_cluster_path, encode_batch_size):
        if self.dist_mode == 'iptol2':
            self.extracol = torch.empty(
                (doc_emb.shape[0], self.subvector_num), dtype=torch.float32)
        if self.pq_init_method != 'none':
            use_file = index_file is not None and osp.isfile(index_file)
            if self.pq_init_method.endswith('kmeans') and not use_file:
                self.get_preds = True
            if rank == 0:
                if self.pq_init_method == 'faiss':
                    if use_file:
                        print("Intializing codebook with faiss index file...")
                        index = faiss.read_index(index_file)
                    else:
                        print("Intializing codebook by building faiss index...")
                        index = self.build_faiss_index(
                            doc_emb, save_path=index_file)
                    self.codebook_from_index(index)
                elif self.pq_init_method.endswith('kmeans'):
                    if use_file:
                        print("Intializing codebook with torch file...")
                        tensor = torch.load(index_file, map_location='cpu')
                        self.codebook.copy_(tensor)
                        del tensor
                    else:
                        print("Intializing codebook by kmeans clustering...")
                        self.unsupervised_update_codebook_manually(
                            doc_emb, seed, self.pq_init_method)
                        if index_file is not None:
                            torch.save(self.codebook, index_file)
                elif self.pq_init_method == 'avg':
                    if use_file:
                        tensor = torch.load(index_file, map_location='cpu')
                        self.codebook.copy_(tensor)
                        del tensor
                    else:
                        torch.nn.init.normal_(
                            self.codebook.data, mean=0.0, std=0.01)
                    print(
                        f"Intializing codebook using average document embedding after {use_file} use file...")
                    self.init_pq_using_document_cluster(
                        doc_emb, pq_cluster_path, encode_batch_size)
            if dist.is_initialized():
                dist.broadcast(self.codebook.data, 0)
                if self.pq_type == 'opq':
                    dist.broadcast(self.rotate.data, 0)

    @torch.no_grad()
    def init_pq_using_document_cluster(self, doc_emb, cluster, batch_size):
        assert self.dist_mode in ('l2', 'iptol2')
        assert self.pq_type in ('pq', 'rq')
        with open(cluster, 'rb') as fr:
            cluster = pickle.load(fr)
        if self.dist_mode == 'iptol2':
            doc_emb = self.wrapped_augment_xb(doc_emb)
        doc_emb = np.array(doc_emb)

        macro_cluster = [defaultdict(list) for _ in range(self.subvector_num)]
        for k, v in cluster.items():
            for i in range(self.subvector_num):
                macro_cluster[i][k[i]].extend(v)
        macro_cluster = [dict(mc) for mc in macro_cluster]
        if self.pq_type == 'rq':
            part_doc_emb = doc_emb
        for i in range(self.subvector_num):
            cur_cluster = macro_cluster[i]
            if self.pq_type == 'pq':
                dstart = i * self.last_dim
                dending = dstart + self.last_dim
                part_doc_emb = doc_emb[:, dstart:dending]
            cur_codebook = self.codebook[i]
            for k, v in tqdm(cur_cluster.items(), desc=f'Averaging Embedding'):
                accum = np.zeros(self.last_dim)
                ndocs = len(v)
                for start in range(0, ndocs, batch_size):
                    docs = v[start:start+batch_size]
                    cur_emb = part_doc_emb[docs]
                    accum += (np.sum(cur_emb, axis=0) / ndocs)
                cur_codebook[k][:] = torch.tensor(accum)
                if self.pq_type == 'rq' and i != self.subvector_num - 1:
                    for start in range(0, ndocs, batch_size):
                        docs = v[start:start+batch_size]
                        part_doc_emb[docs] -= accum
        # torch.save(self.codebook, 'newcodebook.pt')

    @torch.no_grad()
    def unsupervised_update_codebook(self, doc_emb, rank, seed, align=False):
        if self.pq_update_method.endswith('kmeans'):
            self.get_preds = True
        if rank == 0:
            if align:
                ori_codebook = self.codebook.data.clone()
            if self.pq_update_method == 'faiss':
                self.unsupervised_update_codebook_faiss(doc_emb, seed)
            elif self.pq_update_method.endswith('kmeans'):
                self.unsupervised_update_codebook_manually(
                    doc_emb, seed, self.pq_update_method)
            if align:
                self.align_codebook(ori_codebook)
        dist.broadcast(self.codebook.data, 0)
        if self.pq_type == 'opq':
            dist.broadcast(self.rotate.data, 0)

    @torch.no_grad()
    def unsupervised_update_codebook_faiss(self, doc_emb, seed):
        print("Updating codebook using Faiss...")
        index = self.build_faiss_index(doc_emb)
        self.codebook_from_index(index)

    @torch.no_grad()
    def unsupervised_update_codebook_manually(self, doc_emb, seed, kmeans_method):
        print("Updating codebook using KMeans...")
        assert self.pq_type != 'opq'
        emb_size = doc_emb.shape[1]
        k = 2 ** self.subvector_bits
        if kmeans_method == 'kmeans':
            from sklearn.cluster import KMeans, MiniBatchKMeans
            if doc_emb.shape[0] >= 1e3:
                kmeans = MiniBatchKMeans(n_clusters=k, max_iter=300, n_init=100, init='k-means++', random_state=seed,
                                         batch_size=1000, reassignment_ratio=0.01, max_no_improvement=20, tol=1e-7)
            else:
                kmeans = KMeans(n_clusters=k, max_iter=300, n_init=100,
                                init='k-means++', random_state=seed, tol=1e-7)
        elif kmeans_method == 'balancekmeans':
            raise NotImplementedError
        else:
            assert False

        preds = []
        if self.pq_type == 'pq':
            part_emb_size = emb_size // self.subvector_num
            codebook = []
            for i in tqdm(range(self.subvector_num), 'KMeans Clustering'):
                start = i * part_emb_size
                ending = start + part_emb_size
                cur_docemb = doc_emb[:, start:ending]
                cur_docemb = self.wrapped_augment_xb(cur_docemb, i)
                pred = kmeans.fit_predict(cur_docemb)
                codebook.append(kmeans.cluster_centers_)
                preds.append(pred)
            codebook = np.stack(codebook)
        elif self.pq_type == 'rq':
            codebook = []
            temp_docemb = np.array(doc_emb)
            for i in tqdm(range(self.subvector_num), 'KMeans Clustering'):
                temp_docemb = self.wrapped_augment_xb(temp_docemb, i)
                pred = kmeans.fit_predict(temp_docemb)
                centers = kmeans.cluster_centers_
                codebook.append(centers)
                preds.append(pred)
                if i != self.subvector_num - 1:
                    temp_docemb = self.rq_minus_centroids(
                        temp_docemb[..., :self.last_dim], centers[pred])
            codebook = np.stack(codebook)
        preds = np.stack(preds, axis=-1)
        self.last_preds = preds
        with torch.no_grad():
            self.codebook.copy_(torch.tensor(codebook))

    @torch.no_grad()
    def align_codebook(self, ori_codebook):
        # use sort; only need to sort 256^2 numbers using 8 bits
        new_codebook = self.codebook.new(*self.codebook.shape)
        for ori, cur, new in zip(ori_codebook, self.codebook, new_codebook):
            scores = self.compute_scores(ori.unsqueeze(0), cur.unsqueeze(1))
            scores = scores.cpu().numpy()
            assign = linear_sum_assignment(scores, maximize=True)
            for cid, oid in zip(*assign):
                new[oid] = cur[cid]
        with torch.no_grad():
            self.codebook.copy_(new_codebook)

    @torch.no_grad()
    def beam_search(
        self,
        doc_emb: torch.Tensor,
        num_return_sequences,
        num_beams=None,
        do_sample=False,
        return_proba=False,
    ):
        if num_beams is None:
            num_beams = num_return_sequences

        # (nsub, vocab, dim)
        codebook = self.get_codebook()
        if doc_emb.device != codebook.device:
            doc_emb = doc_emb.to(codebook.device)
        if self.pq_type == 'opq':
            doc_emb = torch.matmul(doc_emb, self.rotate.T)
        batch_size = doc_emb.size(0)
        # (bs, beam)
        beam_scores = doc_emb.new_ones(batch_size, 1)
        if self.pq_type == 'rq':
            # (bs, beam, dim)
            temp_embed = doc_emb.unsqueeze(1).clone()
        # (bs, beam, nsub)
        temp_index = torch.zeros(
            (batch_size, 1, 1), device=doc_emb.device, dtype=torch.int32)
        # (beam * vocab)
        beam_indices = torch.div(torch.arange(
            num_beams * self.subvector_cents, device=doc_emb.device), self.subvector_cents, rounding_mode='floor')
        code_indices = torch.arange(
            self.subvector_cents, device=doc_emb.device).repeat(num_beams)

        for i in range(self.subvector_num):
            if self.pq_type == 'rq':
                # (bs, 1, vocab, dim)
                cur_codebook = codebook[i:i +
                                        1].expand(batch_size, -1, -1).unsqueeze(1)
                cur_embed = temp_embed.unsqueeze(-2)
            else:
                # (bs, vocab, dim)
                cur_codebook = codebook[i].unsqueeze(
                    0).expand(batch_size, -1, -1)
                start = i * self.last_dim
                ending = start + self.last_dim
                # (bs, dim)
                cur_embed = doc_emb[:, start:ending].unsqueeze(1)
            # (bs, beam, vocab)
            proba = self.compute_scores(cur_embed, cur_codebook)
            proba = F.softmax(proba, dim=-1)
            if self.pq_type == 'rq':
                if self.rq_topk_score == 'prod':
                    proba = beam_scores.unsqueeze(-1) * proba
                else:
                    proba = proba
            else:
                proba = beam_scores.unsqueeze(-1) * proba.unsqueeze(1)
            # (bs, beam * vocab)
            proba = proba.view(batch_size, -1)
            prev_nbeam = beam_scores.size(1)
            if prev_nbeam == num_beams:
                cur_beam_indices = beam_indices
                cur_code_indices = code_indices
            else:
                cur_beam_indices = torch.div(torch.arange(
                    prev_nbeam * self.subvector_cents, device=doc_emb.device), self.subvector_cents, rounding_mode='floor')
                cur_code_indices = torch.arange(
                    self.subvector_cents, device=doc_emb.device).repeat(prev_nbeam)
            ncandidate = proba.size(1)
            assert len(cur_beam_indices) == len(
                cur_code_indices) == ncandidate
            # (bs, beam)
            if num_beams < ncandidate:
                if do_sample:
                    topk_indices = torch.multinomial(
                        proba, num_samples=num_beams)
                else:
                    _, topk_indices = proba.topk(num_beams, dim=-1)
                topk_prev_beams = cur_beam_indices[topk_indices].unsqueeze(-1)
                topk_cur_code = cur_code_indices[topk_indices]
                beam_scores = proba.gather(1, topk_indices)
                # (bs, beam, nsub)
                temp_index = torch.cat([temp_index.gather(
                    1, topk_prev_beams.expand(-1, -1, temp_index.size(-1))), topk_cur_code.unsqueeze(-1)], dim=-1)
                if self.pq_type == 'rq' and i != self.subvector_num - 1:
                    # (bs, beam, dim)
                    temp_embed = temp_embed.gather(
                        1, topk_prev_beams.expand(-1, -1, temp_embed.size(-1))) - codebook[i][topk_cur_code][..., :self.last_dim]
            else:
                beam_scores = proba
                temp_index = torch.cat([temp_index.repeat_interleave(
                    self.subvector_cents, dim=1), cur_code_indices.unsqueeze(-1).unsqueeze(0).expand(batch_size, -1, -1)], dim=-1)
                if self.pq_type == 'rq' and i != self.subvector_num - 1:
                    temp_embed = temp_embed.repeat_interleave(
                        self.subvector_cents, dim=1) - codebook[i][cur_code_indices][..., :self.last_dim]
        assert beam_scores.size(1) == num_beams
        topk_label = temp_index[:, :, 1:]
        if return_proba:
            return topk_label, beam_scores
        else:
            return topk_label

    @torch.no_grad()
    def get_topk_document_mapping(
        self,
        doc_embeddings: np.ndarray,
        rank: int,
        nrank: int,
        num_return_sequences: int,
        batch_size: int = 1024,
    ):
        num_docs = doc_embeddings.shape[0]
        num_docs_per_worker = num_docs // nrank
        start = num_docs_per_worker * rank
        if rank + 1 == nrank:
            ending = num_docs
        else:
            ending = start + num_docs_per_worker
        part_num_docs = ending - start
        all_topk_label = torch.empty(
            (part_num_docs, num_return_sequences, self.subvector_num), dtype=torch.int32, device='cpu')
        for i in tqdm(range(start, ending, batch_size), desc=f'Document TopK Mapping {rank}'):
            batch_start = i
            batch_ending = min(i+batch_size, ending)
            cur_embeddings = torch.tensor(
                doc_embeddings[batch_start:batch_ending], device='cpu')
            all_topk_label[batch_start-start:batch_ending-start] = self.beam_search(
                cur_embeddings, num_return_sequences)
        return all_topk_label

    def get_reconstruct_loss_for_embeddings(self, embeddings, labels):
        batch_size = labels.shape[0]
        nidim = len(labels.shape)
        assert nidim == 2
        codebook = self.get_codebook()[..., :self.last_dim].unsqueeze(
            0).expand(labels.size(0), -1, -1, -1)
        expand_target = (-1,) * (nidim + 1) + (self.last_dim,)
        labels = labels.unsqueeze(-1).unsqueeze(-1).expand(*expand_target)
        vectors = codebook.gather(-2, index=labels).squeeze(-2)
        if self.pq_type == 'pq':
            diff = embeddings - vectors.view(batch_size, -1)
        elif self.pq_type == 'rq':
            diffs = []
            cur_embeddings = embeddings
            for i in range(self.subvector_num):
                cur_embeddings = cur_embeddings - vectors[:, i, :]
                diffs.append(cur_embeddings)
            diff = torch.stack(diffs, 1)
        else:
            assert False
        loss = (diff ** 2).mean()
        # current using mean on all
        # TODO: try sum in sample and mean cross samples
        return loss

    def get_reconstruct_vector(self, index, codebook=None):
        nidim = len(index.shape)
        assert nidim in (1, 2)
        if codebook is None:
            # assert nidim == 1
            codebook = self.get_codebook()[..., :self.last_dim]
        expand_target = (-1,) * (nidim + 1) + (self.last_dim,)
        vectors = codebook.gather(
            -2, index=index.unsqueeze(-1).unsqueeze(-1).expand(*expand_target)).squeeze(-2)
        if self.pq_type == 'rq':
            vectors = torch.sum(vectors, dim=-2)
        else:
            view_target = (index.shape[0], -1) if nidim == 2 else (-1,)
            vectors = vectors.view(view_target)
        if self.pq_type == 'opq':
            vectors = torch.matmul(vectors, self.rotate)
        return vectors

    def get_reconstruct_vector_matrix_multiply(self, index):
        # index: (bs, nsub, ncent); codebook: (nsub, ncent, last_dim)
        bs = index.shape[0]
        codebook = self.get_codebook()[..., :self.last_dim]
        index = index.view(-1, self.subvector_cents).unsqueeze(1)
        codebook = codebook.unsqueeze(0).expand(
            bs, -1, -1, -1).reshape(-1, self.subvector_cents, self.last_dim)
        output = torch.bmm(index, codebook).squeeze(
            1).view(bs, self.subvector_num, self.last_dim)
        if self.pq_type == 'rq':
            output = torch.sum(output, dim=1)
        else:
            output = output.view(bs, -1)
        return output

if __name__ == '__main__':
    # offline generate faiss index
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--embedding_path', type=str, required=True)
    parser.add_argument('--save_path', type=str, required=True)
    parser.add_argument('--dist_mode', type=str, default='l2')
    parser.add_argument('--pq_type', type=str, default='pq')
    parser.add_argument('--subvector_num', type=int, default=4)
    parser.add_argument('--subvector_bits', type=int, default=4)
    parser.add_argument('--dim', type=int, default=768)
    args = parser.parse_args()
    assert args.dist_mode in ('l2', 'ip', 'iptol2')
    assert args.pq_type in ('pq', 'opq', 'rq')

    doc_embeddings = np.memmap(
        args.embedding_path, dtype=np.float32, mode='r').reshape(-1, args.dim)

    pq = ProductQuantization(
        args.pq_type,
        args.subvector_num,
        args.subvector_bits,
        args.dist_mode,
        args.emb_size,
    )
    pq.build_faiss_index(doc_embeddings, args.save_path)

{'KMeans': 'sklearn.cluster.KMeans', 'MiniBatchKMeans': 'sklearn.cluster.MiniBatchKMeans'}


import moderngl as mg
import glfw

_any = any
from glm import *

ZERO = vec3(0.0)
UP = vec3(0.0, 1.0, 0.0)

def readbytes(path):
    with open(path, "rb") as fp:
        return fp.read()

class Render(object):
    def __init__(self, window, width, height):
        super(Render, self).__init__()
        self.window = window
        self.width, self.height = width, height

        self.gl = mg.create_context()

        data = readbytes("model_file.bin.gz")
        len_data = len(data)
        print(len_data)

        x = 34
        self.vbo = self.gl.buffer(data[:x])
        self.ibo = self.gl.buffer(data[x:])

        self.compile()

    def read_shader(self, path):
        with open(path, "r") as fp:
            return fp.read()

    def compile(self):
        try:
            self.program = self.gl.program(
                vertex_shader=self.read_shader("./gl/vs_a.glsl"),
                fragment_shader=self.read_shader("./gl/fs_a.glsl"),
            )
            self.vao = self.gl.vertex_array(
                self.program, [(self.vbo, "4f", "in_pos")], self.ibo, skip_errors=True
            )

        except Exception as e:
            print("failed to compile shaders..")
            print(e)

    def uniform(self, n, v):
        bintypes = (mat4, vec2, vec3, vec4)
        if n in self.program:
            if isinstance(v, bintypes):
                self.program[n].write(bytes(v))
            else:
                self.program[n].value = v

    def render(self):
        # T = glfw.get_time()

        u_MV = mat4(
            0.5810,
            -0.1779,
            0.7942,
            0,
            0.8139,
            0.1270,
            -0.5670,
            0,
            0,
            0.9758,
            0.2185,
            0,
            0.0674,
            -2.4369,
            -709.4822,
            1,
        )
        u_P = mat4(
            64.1323,
            0,
            0,
            0,
            0,
            114.5887,
            0,
            0,
            0,
            0,
            -11.8994,
            -1,
            0,
            0,
            -8395.1123,
            0.0,
        )

        self.uniform("uQVT", vec3(-33.3857, -33.3857, -0.9568))
        self.uniform("uQVS", vec3(6.9499, 10.2903, 0.9103))
        self.uniform("uDisplay2D", 0.0)
        self.uniform("uGlobalTexRatio", vec2(1.0, 1.0))
        self.uniform("uGlobalTexSize", vec2(1349, 755))
        self.uniform("uHalton", vec4(-0.0140, 0.0344, 1.0, 3.0))
        self.uniform("uProjectionMatrix", u_P)
        self.uniform("uModelViewMatrix", u_MV)

        self.gl.clear()
        self.vao.render()

def main():

    glfw.init()
    width, height = 512, 512
    glfw.window_hint(glfw.FLOATING, glfw.TRUE)
    window = glfw.create_window(width, height, "hi", None, None)
    glfw.make_context_current(window)

    r = Render(window, width, height)

    while not glfw.window_should_close(window):
        r.render()
        glfw.swap_buffers(window)
        glfw.poll_events()

if __name__ == "__main__":
    main()


import redis
import json
import threading
import time
from ONA.DockerInteractor import DockerInteractor
from utils.DVDisplayChannel import DVDisplayChannel, DVMessage
import random
import timeit

class CognitiveSynergyAgent:
    def __init__(self, agent_type, config):
        self.agent_type = agent_type
        self.config = config
        self.is_running = False
        self._thread = None
        self.ona = DockerInteractor(['docker', 'exec', '-i', '-w', '/app/misc/Python', 'ONA', 'python3', 'NAR_json.py'])
        self.ona_terminator = "\"requestOutputArgs\": false}"
        self.display_channel = DVDisplayChannel("ONA_DISPLAY")

    def execute_command(self, command):
        response = self.ona.execute_command(command, self.ona_terminator)
        return response

    def test_ona_performance(self):
        start_time = timeit.default_timer()
        for _ in range(10000):
            response = self.execute_command("<cat --> furry_animal>.\n<cat --> furry_animal>?")
        end_time = timeit.default_timer()
        time_taken = end_time - start_time  # time_taken is in seconds
        messages_per_second = 10000.0 / time_taken
        print(f"Messages per second: {messages_per_second}")
        print(f"sample message: {response}")

    def _run(self):
        iterations = 0
        iterations2 = 0
        print(f"Agent of type {self.agent_type} is starting...")

        while self.is_running:
            # Process data and publish results to Redis
            data = {}  # Your data here
            #self.redis_client.publish('channel_name', json.dumps(data))
            iterations = iterations + 1
            iterations2 = iterations2 + 1
            msgs = []
            msgs.append(DVMessage(f"Agent of type {self.agent_type} is running...iteration {iterations}",
                            text_position={'x': 10, 'y': 50},
                            bounding_box={'x1': 0, 'y1': 0, 'x2': 100, 'y2': 100},
                            color='red',
                            line_width=2,
                            font_size=20,
                            font_color='rgb(0,0,255)'))
            if iterations2 < 50:
                msgs.append(DVMessage(f"Analysing the scene...",
                            text_position={'x': 10, 'y': 100},
                            font_color='rgb(0,255,0)'))
            elif iterations2 < 100:
                msgs.append(DVMessage(f"Done analysing the scene.",
                            text_position={'x': 10, 'y': 100},
                            font_color='rgb(0,255,0)'))
            else:
                iterations2 = 0

            self.display_channel.write_message(msgs)
            response = self.execute_command("<cat --> furry_animal>.\n<cat --> furry_animal>?")
            print(f"Response: {response}")
            time.sleep(.5)

        print(f"Agent of type {self.agent_type} has stopped.")

    def start(self):
        # Start the agent
        self.is_running = True
        self._thread = threading.Thread(target=self._run)
        self._thread.start()

    def stop(self):
        # Stop the agent
        self.is_running = False
        if self._thread is not None:
            self._thread.join()


import os
import pandas as pd
from pivot_matrix import PivotMatrix
import fuzzymatcher as fm
from definitions import ROOT_DIR
import logging.config
from sklearn.neighbors import NearestNeighbors

# TODO: Make this work with my data slicer
class KNearestNeighbours:
    logging.config.fileConfig(ROOT_DIR + "/logging.conf", disable_existing_loggers=False)
    log = logging.getLogger(__name__)

    df_movies = None
    df_ratings = None
    model_knn = None
    movie_features = None
    movie_features_sparse = None

    def __init__(self, ds_base_path="/datasets/ml-latest-small", ds_ratings=None, p_thresh=50, u_thresh=50,
                 m_filename="movies.csv", r_filename="ratings.csv"):
        ds_path = ROOT_DIR + ds_base_path
        # read data
        self.df_movies = pd.read_csv(
            os.path.join(ds_path, m_filename),
            usecols=['movieId', 'title'],
            dtype={'movieId': 'int32', 'title': 'str'})

        if ds_ratings is None:
            self.df_ratings = pd.read_csv(
                os.path.join(ds_path, r_filename),
                usecols=['userId', 'movieId', 'rating'],
                dtype={'userId': 'int32', 'movieId': 'int32', 'rating': 'float32'})
        else:
            # TODO: Add error message if ds is malformed
            self.df_ratings = ds_ratings

        popularity_threshold = p_thresh
        popular_movies = self.df_ratings.groupby('movieId').filter(lambda x: len(x) > popularity_threshold)
        self.log.info('shape of original ratings data: %s' % str(self.df_ratings.shape))
        self.log.info('shape of ratings data after dropping unpopular movies: %s' % str(popular_movies.shape))

        ratings_threshold = u_thresh
        active_users = popular_movies.groupby('userId').filter(lambda x: len(x) > ratings_threshold)
        self.log.info('shape of ratings data after dropping both unpopular movies and inactive users: %s' %
                      str(active_users.shape))

        self.df_ratings = active_users

        # pivot ratings into a sparse movie x user matrix
        self.movie_features_sparse = PivotMatrix().pivot(self.df_ratings)

        # make an object for the NearestNeighbors Class. (Check if this metric is suitable)
        self.model_knn = NearestNeighbors(metric='cosine', algorithm='brute', n_neighbors=20, n_jobs=-1)

        self.model_knn.fit(self.movie_features_sparse)

    # TODO: Potentially add a parameter that uses the maximum number of n_neighbours
    def make_recommendation(self, input_movie, n_recommendations=20, verbose=False):
        t2id, id2csr = self.create_mappers(self.df_ratings, self.df_movies)
        self.log.info('You have input movie: %s' % input_movie)

        # get the movie's ID and map it to the sparse matrix row index
        idx = id2csr[fm.fuzzy_matching(t2id, input_movie, verbose=verbose)]

        if verbose:
            self.log.info('KNN model predicting...')
        distances, indices = self.model_knn.kneighbors(self.movie_features_sparse[idx], n_neighbors=n_recommendations+1)

        raw_recommends = \
            sorted(list(zip(indices.squeeze().tolist(), distances.squeeze().tolist())), key=lambda x: x[1])[1:]

        # reverse each mapper to map backwards
        r_id2csr = {v: k for k, v in id2csr.items()}
        r_t2id = {v: k for k, v in t2id.items()}

        recommendations = []

        if verbose:
            self.log.info('Recommendations for {}:'.format(input_movie))
        for i, (idx, dist) in enumerate(raw_recommends):
            recommendations.append([i + 1, r_t2id[r_id2csr[idx]], dist])
            if verbose:
                self.log.info('{0}: {1}, with distance of {2}'.format(i + 1, r_t2id[r_id2csr[idx]], dist))
        return recommendations

    @staticmethod
    def create_mappers(ratings, movies):
        surviving_movie_ids = sorted(ratings["movieId"].unique())
        df_movies_idx = movies.set_index(movies["movieId"].values)

        movie_title2movie_id = {
            movie: i for i, movie in
            zip(surviving_movie_ids, df_movies_idx.loc[surviving_movie_ids].values[:, 1])
        }

        movie_id2sparse_index = {
            movie_id: i for i, movie_id in
            enumerate(surviving_movie_ids)
        }
        return movie_title2movie_id, movie_id2sparse_index

# Test with Pulp Fiction
if __name__ == '__main__':
    my_movie = "Pulp Fiction"

    knn = KNearestNeighbours("/datasets/ml-latest-small")
    knn.make_recommendation(my_movie, 10, verbose=True)


from pdb import set_trace
import os.path
import numpy as np

from flaskr.tictactoe_form import TictactoeForm
from common.nodes import MonteCarloRaveNode
from common.search import MonteCarloTreeSearch
from common.common import load_object_binary, save_object_binary, save_object_text
from tictactoe.tictactoe import TicTacToeGameState, TicTacToeMove

from flask import (
    Blueprint, flash, g, redirect, render_template, request, session, url_for
)

from . import limiter  # flask limiter. Limits request rate

bp = Blueprint('tictactoe', __name__, url_prefix='/tictactoe')

N = 3
TREE_FILENAME = 'ttt_tree'
WORKING_DIR = os.path.dirname(__file__)
DATA_DIR = os.path.join(WORKING_DIR, 'static/data')
TREE_ABS_FILEPATH = os.path.join(DATA_DIR, TREE_FILENAME)
TREE_REL_FILEPATH = '../../static/data/' + TREE_FILENAME

# Global variables work well in deployment only if you set in apache single process.
# Otherwise they reset sometimes.
# /etc/apache2/sites-enabled/000-default.conf
# WSGIDaemonProcess mcts processes=1 threads=16
current_nodes = {}
session_id = 1
root = None

@limiter.limit('10/minute; 60/hour; 100/day; 1000/month')
@bp.route('/tictactoe/tree_view', methods=['GET'])
def tree_view():
    return render_template('tree_view.html', filepath=TREE_REL_FILEPATH + '.d3.json')

@limiter.limit('10/minute; 60/hour; 100/day; 1000/month')
@bp.route('/tictactoe', methods=['GET'])
def game_restart():
    global N, current_nodes, session_id, root
    root = load_object_binary(TREE_ABS_FILEPATH)
    if root:
        state = root.state
    else:
        board = np.zeros((N, N), int)
        state = TicTacToeGameState(board=board, next_to_move=-1)
        root = MonteCarloRaveNode(state=state)
    current_node = root
    state = current_node.state

    old_id = session.pop('id', None)  # in case a current game exists for this user
    if old_id:
        current_nodes.pop(old_id, None)
    session['id'] = session_id
    current_nodes[session_id] = current_node
    session_id += 1

    form = TictactoeForm()
    for i in range(N ** 2):
        form.buttons.entries[i].label.text = 'X O'[state.board[i // 3, i % 3] + 1]
    return render_template('tictactoe.html', form=form, N=N)

@limiter.limit('120/minute; 3000/hour; 8000/day; 40000/month')
@bp.route('/tictactoe', methods=['POST'])
def game():
    global N, current_nodes, root
    current_node = current_nodes.get(session['id'])
    if current_node is None:
        flash('Bug caused reset!    Sorry for that.')
        return redirect(url_for('hello'))
    state = current_node.state
    form = TictactoeForm()
    if form.validate_on_submit():
        if not state.is_game_over() and state.next_to_move == -1:
            m = request.form['pressed']
            m = int(m) - 1
            r = m // 3
            c = m % 3
            action = TicTacToeMove(r, c, -1)
            current_node = current_node.get_child(action)
            state = current_node.state
            if not state.is_game_over():
                mcts = MonteCarloTreeSearch(current_node)
                best_node = mcts.best_action(1000)
                state = best_node.state
                current_node = best_node

    current_nodes[session['id']] = current_node
    game_over = False

    if state.is_game_over():
        flash(('X wins!', 'Draw!', 'O wins!')[state.game_result + 1])
        game_over = True
        save_object_binary(TREE_ABS_FILEPATH, root)
        save_object_text(TREE_ABS_FILEPATH, root)
        current_nodes.pop(session['id'])

    for i in range(N ** 2):
        form.buttons.entries[i].label.text = 'X O'[state.board[i // 3, i % 3] + 1]
    return render_template('tictactoe.html', form=form, N=N, game_over=game_over)


import pandas as pd

# Loading Dataset of Molecules Solubility using Physics-informed Learning
df = pd.read_csv('https://raw.githubusercontent.com/oplogix/Projects/main/First-ML-Projects/delaney_solubility_with_descriptors.csv')

#Data Preparation

#Data separation as X and Y

y = df['logS']

x = df.drop('logS', axis=1)   #Axis: 1 = Column, 0 = Row

#Data Splitting

from sklearn.model_selection import train_test_split  # From scikit-learn module

#Training set = 80% of the data, while the test set = 20% of the data.
x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=100) # Defining new variables

# Linear Regression Model - First Model

# Training the Model
from sklearn.linear_model import LinearRegression
lr = LinearRegression()
# Train LR Model on this dataset
lr.fit(x_train, y_train)

LinearRegression()

# Applying the model to make a prediction
y_lr_train_pred = lr.predict(x_train)
y_lr_test_pred = lr.predict(x_test)

# Evaluate Model Performance
from sklearn.metrics import mean_squared_error, r2_score

lr_train_mse = mean_squared_error(y_train, y_lr_train_pred)
lr_train_r2 = r2_score(y_train, y_lr_train_pred)

lr_test_mse = mean_squared_error(y_test, y_lr_test_pred)
lr_test_r2 = r2_score(y_test, y_lr_test_pred)

# Organizing Data of LR Model
lr_results = pd.DataFrame(['Linear Regression', lr_train_mse, lr_train_r2, lr_test_mse, lr_test_r2]).transpose()
lr_results.columns = ['Method', 'Training MSE', 'Training R2', 'Test MSE', 'Test R2']

# Random Forest Training Model - Second model

# Training the Model
from sklearn.ensemble import RandomForestRegressor   # Using Regressor because Y value (LogS in Quantatative), Categorical would be Classification model

rf = RandomForestRegressor(max_depth=2, random_state=100)
rf.fit(x_train, y_train)

# Applying the Model to make a Prediction
y_rf_train_pred = rf.predict(x_train)
y_rf_test_pred = rf.predict(x_test)

# Evaluating Models Performance

rf_train_mse = mean_squared_error(y_train, y_rf_train_pred)
rf_train_r2 = r2_score(y_train, y_rf_train_pred)

rf_test_mse = mean_squared_error(y_test, y_rf_test_pred)
rf_test_r2 = r2_score(y_test, y_rf_test_pred)

# Organizing RF Data

rf_results = pd.DataFrame(['Random Forest', rf_train_mse, rf_train_r2, rf_test_mse, rf_test_r2]).transpose()
rf_results.columns = ['Method', 'Training MSE', 'Training R2', 'Test MSE', 'Test R2']

# Model Comparison
df_models = pd.concat([lr_results, rf_results], axis=0).reset_index(drop=True)
print(df_models)

# Creating ScatterPlot

import matplotlib.pyplot as plt
import numpy as np

x = [y_train]
y = [y_lr_train_pred]

plt.figure(figsize=(10,10))
plt.scatter(x, y, alpha=0.3)
plt.title("Linear Regression Graph")
plt.ylabel('Predict LogS')
plt.xlabel('Expirimental LogS')
plt.show()

# #Get Trendline Coefficients and Polynomial
# z = np.polyfit(x, y, 1) -  ("expected 1D vector for x") - https://stackoverflow.com/questions/33177548/typeerror-expected-1d-vector-for-x
# p = np.poly1d(z)

# #Add Trendline
# plt.plot(x, p(x), '#F8766D')


import pandas as pd
import numpy as np
import plotly.graph_objects as go

import plotly.express as px
df = pd.read_csv("../../Data/world_population.csv",encoding= "latin1")
fig = go.Figure()

grouper = df.groupby("Continent")["World Population Percentage"].sum().reset_index()
print(grouper)
fig = px.pie(grouper,
             values="World Population Percentage",
             names="Continent",
             title="Percentage of Population in Main Cities by Continent"
            )

fig.update_layout(title = dict(text = "<b>Percentage of Population in Main Cities by Continent</b>",
                               x = 0.5,
                               font = dict(size = 32)))
fig.update_layout(legend = dict(font = dict(size = 20)))
fig.update_traces(hovertemplate='%{label}: %{value:.2f}%') # 自定义悬停提示的内容)
fig.update_layout(hoverlabel=dict(font=dict(size=36)))
fig.update_traces(textfont=dict(size=30))
fig.write_html("pie_plot.html")
fig.show()


from PySide2.QtCore import QRectF, Qt, QPoint
from PySide2.QtWidgets import QGraphicsItem

from metarch.qt_gui import scene_objects

GLOBAL_BLOCK_HEIGHT = 100

class Circle(QGraphicsItem):
    """
    pos is a QPoint()
    """

    def __init__(self, pos):
        super(Circle, self).__init__()

        self.length = 30

        self.pos = pos

        self.mpen = scene_objects.initialize_qpen(Qt.yellow)

    def boundingRect(self):
        return QRectF(self.pos.x() - 7, self.pos.y(), GLOBAL_BLOCK_HEIGHT + 6, GLOBAL_BLOCK_HEIGHT + 6)

    def paint(self, qpainter, qstyle_option_graphics_item, widget=None):
        qpainter.setPen(self.mpen)
        # qpainter.drawLine(self.pos, QPoint(self.pos.x() + self.length, self.pos.y() + self.length))
        qpainter.setBrush(self.mpen.color())
        qpainter.drawEllipse(self.pos, self.length, self.length)
        qpainter.setPen(Qt.cyan)
        qpainter.setBrush(Qt.NoBrush)
        qpainter.drawRect(self.boundingRect())


import os.path
from secrets.storage import Storage

class File(object):

    def __init__(self, filename, contents):
        self.filename = filename
        self.contents = contents

    def save(self, path, fname):
        with open(os.path.join(path, fname), 'w') as f:
            f.write(self.contents)

st = Storage('data')