import torch
from torch.autograd import Variable
import onmt.io
from onmt.Utils import aeq, use_gpu

from cocoa.core.entity import is_entity
from cocoa.neural.generator import Generator, Sampler

from .symbols import markers, category_markers, sequence_markers
from .utterance import UtteranceBuilder

from cocoa.core.entity import is_entity, Entity, CanonicalEntity

import numpy as np

class LFSampler(Sampler):

    var_for_price = 0.1
    INTENT_NUM = -1

    def __init__(self, model, vocab,
                 temperature=1, max_length=100, cuda=False, model_type='sl'):
        super(LFSampler, self).__init__(model, vocab, temperature=temperature, max_length=max_length, cuda=cuda)
        self.acc_or_rej = list(map(self.vocab.to_ind, ('accept', 'reject')))
        self.offer = list(map(self.vocab.to_ind, ('offer', )))
        self.price_actions = list(map(self.vocab.to_ind, ('counter', 'propose', 'offer', 'agree', 'disagree')))
        self.model_type = model_type

        self.policy_history = []

        # init rl actions
        _ = self.rl_actions
        LFSampler.INTENT_NUM = len(self.vocab.ind_to_word)

    @staticmethod
    def softmax_with_mask(policy, mask=1):
        if isinstance(mask, torch.Tensor):
            mask = mask.to(device=policy.device)

        policy.sub_(policy.max(1, keepdim=True)[0].expand(-1, policy.size(1)))
        # mask = batch.policy_mask
        # policy[mask == 0] = -100.
        # print(batch.policy_mask)

        # Avoid policy equal to zero ( + 1e-6 )
        p_exp = (policy.exp() + 1e-6).mul(mask)
        # p_exp = (policy.exp() + 1e-6)
        # if torch.sum(p_exp).item() == 0:
        #     p_exp += torch.ones_like(p_exp).mul(batch.policy_mask)

        policy = p_exp / (torch.sum(p_exp, keepdim=True, dim=1))
        if torch.any(policy < 0) or torch.any(torch.isnan(policy)):
            print('lots of errors: ', p_exp, mask)

        act = torch.multinomial(policy, 1).reshape(-1)

        # (batch_size,), (batch_size, intent_size)
        return act, policy

    def _run_batch_tom_identity(self, batch, hidden_state, only_identity=False, id_gt=False):
        if id_gt:
            id_gt = batch.strategy
        else:
            id_gt = None
        if only_identity:
            identity, next_hidden = \
                self.model.encoder.identity(batch.identity_state, batch.extra, hidden_state, uttr=batch.uttr)
            predictions = None
        else:
            output = self.model(batch.uttr, batch.identity_state, batch.state,
                              batch.extra, hidden_state, id_gt)
            if len(output) == 3:
                predictions, next_hidden, identity = output
            else:
                predictions, next_hidden = output
                identity = None
        return predictions, next_hidden, identity

    def generate_batch(self, batch, gt_prefix=1, enc_state=None, whole_policy=False, special_actions=None,
                       temperature=1, acpt_range=None, hidden_state=None):
        # This is to ensure we can stop at EOS for stateful models
        assert batch.size == 1.

        # Run the model
        if self.model_type in ['sl', 'rl']:
            policy, p_policy = self.model(batch.uttr, batch.state)
            next_hidden = None
        else:
            # tom model
            out, next_hidden, _ = self._run_batch_tom_identity(batch, hidden_state)
            policy, p_policy = out

        last_intent, last_price = batch.get_pre_info(self.vocab)

        p_mask = batch.policy_mask(self.vocab)
        intent, policy = self.softmax_with_mask(policy, p_mask)

        prob = policy[0, intent.item()].item()
        if self.model_type == 'rl':
            p_act, p_policy = self.softmax_with_mask(p_policy)
            price = p_act.item()
            prob = prob*p_policy[0, price].item()
        else:
            price = p_policy.item()

        # Use Normal distribution with constant variance as policy on price

        # Use rule for Supervised learning agent
        if acpt_range is not None:
            assert len(acpt_range) == 2
            # Check if last action is offer
            if last_intent[0] in self.offer:
                offer_price = last_price[0, 1].item()
                policy = policy * 0

                if (acpt_range[0] <= offer_price) and (offer_price <= acpt_range[1]):
                    # Accept
                    act_idx = self.acc_or_rej[0]
                else:
                    act_idx = self.acc_or_rej[1]

                intent[0] = act_idx
                policy[0, act_idx] = 1
                prob = 1.

        # TODO: Not correct, for multiple data.
        original_price = price
        if intent not in self.price_actions:
            price = None

        # print('gen output: ',policy, price)
        ret = {"intent": intent,
               "price": price,
               "policy": policy,
               "p_policy": p_policy,
               'prob': prob,
               'original_price': original_price,
               'rnn_hidden': next_hidden
               }
        # ret["batch"] = batch
        return ret

    def get_sl_actions(self, mean, logstd, sample_num=1):
        # get all actions
        all_actions = []
        num_actions = len(self.vocab.word_to_ind)
        actions = list(range(num_actions))
        for i in actions:
            if not i in self.price_actions:
                all_actions.append([i, None])
            else:
                for j in range(sample_num):
                    all_actions.append([i, mean])
        return all_actions

    _rl_actions = None
    _rl_act_index = None
    PACT_NUM = 4

    @property
    def rl_act_index(self):
        if self._rl_act_index is None:
            _ = self.rl_actions
        return self._rl_act_index

    @property
    def rl_actions(self):
        if LFSampler._rl_actions is not None:
            return LFSampler._rl_actions

        # get all actions
        all_actions = []
        num_actions = len(self.vocab.word_to_ind)
        actions = list(range(num_actions))
        idx = 0
        LFSampler._rl_actions = {}
        for i in actions:
            if not i in self.price_actions:
                all_actions.append((i, None))
                LFSampler._rl_actions[all_actions[-1]] = idx
                idx += 1
            else:
                for j in range(self.PACT_NUM):
                    all_actions.append((i, j))
                    LFSampler._rl_actions[all_actions[-1]] = idx
                    idx += 1

        LFSampler._rl_actions = all_actions
        return all_actions

    def _old_get_all_actions(self, mean, logstd, p_num=5, no_sample=False):
        # prices
        self.prices = []
        step = 0.05
        now = step
        while now < 1:
            self.prices.append(now)
            now += step
        num_acions = len(self.vocab.word_to_ind)
        self.actions = list(range(num_acions))

        # get all actions
        all_actions = []
        for i in self.actions:
            if not i in self.price_actions:
                # print('a ', self.vocab.to_word(i))
                all_actions.append((i,None))
        d = torch.distributions.normal.Normal(mean, logstd.exp())

        for i in self.price_actions:
            # print('pa ', self.vocab.to_word(i))
            if no_sample:
                now = 0.
                step = 0.01
                p_list = []
                while now <= 1:
                    p_list.append(now)
                    now += step
                prob = d.log_prob(torch.from_numpy(np.array(p_list)).float().to(mean.device)).exp()
                prob = prob / prob.sum(dim=-1, keepdim= True)
                for j, p in enumerate(p_list):
                    pp = prob[j].cpu().item()
                    if pp >= 1e-6:
                        all_actions.append((i,p, pp))

            else:
                for j in range(p_num):
                    p = mean + logstd.exp() * torch.randn_like(mean)
                    p = p.cpu().item()
                    all_actions.append((i,p))
        return all_actions

    def get_policyHistogram(self):
        import numpy as np
        import matplotlib.pyplot as plt
        import seaborn as sns
        import re

        allNum = len(self.policy_history)
        tmpData = np.mean(self.policy_history,axis=0)[0]
        r = re.compile(u'\d+[.]?\d*')
        x, w = [], []
        for i in range(len(tmpData)):
            tmp = self.vocab.ind_to_word[i]
            if not is_entity(tmp):
                continue
            name = tmp.canonical.value
            if abs(name) > 10.1: continue
            x.append(name)
            w.append(tmpData[i])

        w = w/np.sum(w)
        from scipy.stats import norm
        sns.distplot(x, bins=100, kde=False, hist_kws={'weights': w}, )
        #plt.show()

def get_generator(model, vocab, scorer, args, model_args):
    from onmt.Utils import use_gpu
    if args.sample:
        if model_args.model == 'lf2lf':
            generator = LFSampler(model, vocab, args.temperature,
                                max_length=args.max_length,
                                cuda=use_gpu(args))
        else:
            generator = Sampler(model, vocab, args.temperature,
                                max_length=args.max_length,
                                cuda=use_gpu(args))
    else:
        generator = Generator(model, vocab,
                              beam_size=args.beam_size,
                              n_best=args.n_best,
                              max_length=args.max_length,
                              global_scorer=scorer,
                              cuda=use_gpu(args),
                              min_length=args.min_length)
    return generator


import numpy as np
import networkx as nx
from sklearn import preprocessing
import time
from tqdm import tqdm
from .. import BaseModel, register_model, alias_draw, alias_setup

@register_model("pte")
class PTE(BaseModel):
    r"""The PTE model from the `"PTE: Predictive Text Embedding through Large-scale Heterogeneous Text Networks"
    <https://arxiv.org/abs/1508.00200>`_ paper.

    Args:
        hidden_size (int) : The dimension of node representation.
        walk_length (int) : The walk length.
        walk_num (int) : The number of walks to sample for each node.
        negative (int) : The number of nagative samples for each edge.
        batch_size (int) : The batch size of training in PTE.
        alpha (float) : The initial learning rate of SGD.
    """

    @staticmethod
    def add_args(parser):
        """Add model-specific arguments to the parser."""
        # fmt: off
        parser.add_argument('--walk-length', type=int, default=80,
                            help='Length of walk per source. Default is 80.')
        parser.add_argument('--walk-num', type=int, default=20,
                            help='Number of walks per source. Default is 20.')
        parser.add_argument('--negative', type=int, default=5,
                            help='Number of negative node in sampling. Default is 5.')
        parser.add_argument('--batch-size', type=int, default=1000,
                            help='Batch size in SGD training process. Default is 1000.')
        parser.add_argument('--alpha', type=float, default=0.025,
                            help='Initial learning rate of SGD. Default is 0.025.')
        # fmt: on

    @classmethod
    def build_model_from_args(cls, args):
        return cls(
            args.hidden_size,
            args.walk_length,
            args.walk_num,
            args.negative,
            args.batch_size,
            args.alpha,
        )

    def __init__(self, dimension, walk_length, walk_num, negative, batch_size, alpha):
        super(PTE, self).__init__()
        self.dimension = dimension
        self.walk_length = walk_length
        self.walk_num = walk_num
        self.negative = negative
        self.batch_size = batch_size
        self.init_alpha = alpha

    def train(self, G, node_type):
        self.G = G
        self.node_type = node_type
        self.num_node = G.number_of_nodes()
        self.num_edge = G.number_of_edges()
        self.num_sampling_edge = self.walk_length * self.walk_num * self.num_node

        self.num_node_type = len(set(self.node_type))
        context_node = [nid for nid, ntype in enumerate(self.node_type) if ntype == 0]

        self.edges, self.edges_prob = [[] for _ in range(self.num_node_type)], []
        self.node_prob, self.id2node = [], [dict() for _ in range(self.num_node_type)]

        subgraphs = []
        for i in range(self.num_node_type):
            for j in range(i + 1, self.num_node_type):
                context_node = [nid for nid, ntype in enumerate(self.node_type) if ntype == i or ntype == j]
                sub_graph = nx.Graph()
                sub_graph = self.G.subgraph(context_node)
                if sub_graph.number_of_edges() != 0:
                    subgraphs.append(sub_graph)
        self.num_graph = len(subgraphs)
        print("number of subgraph", self.num_graph)

        for i in range(self.num_graph):
            self.edges[i] = [[e[0], e[1]] for e in subgraphs[i].edges()]
            edges_prob = np.asarray([subgraphs[i][u][v].get("weight", 1.0) for u, v in self.edges[i]])
            edges_prob /= np.sum(edges_prob)
            edges_table_prob = alias_setup(edges_prob)
            self.edges_prob.append(edges_table_prob)

            context_node = subgraphs[i].nodes()
            self.id2node[i] = dict(zip(range(len(context_node)), context_node))
            node2id = dict(zip(context_node, range(len(context_node))))

            degree_weight = np.asarray([0] * len(context_node))
            for u in context_node:
                for v in list(subgraphs[i].neighbors(u)):
                    degree_weight[node2id[u]] += subgraphs[i][u][v].get("weight", 1.0)

            node_prob = np.power(degree_weight, 0.75)
            node_prob /= np.sum(node_prob)
            nodes_table_prob = alias_setup(node_prob)
            self.node_prob.append(nodes_table_prob)

        print("train pte with 2-order")
        self.emb_vertex = (np.random.random((self.num_node, self.dimension)) - 0.5) / self.dimension
        self.emb_context = self.emb_vertex
        self._train_line()
        embedding = preprocessing.normalize(self.emb_vertex, "l2")
        return embedding

    def _update(self, vec_u, vec_v, vec_error, label):
        # update vetex embedding and vec_error
        f = 1 / (1 + np.exp(-np.sum(vec_u * vec_v, axis=1)))
        g = (self.alpha * (label - f)).reshape((len(label), 1))
        vec_error += g * vec_v
        vec_v += g * vec_u

    def _train_line(
        self,
    ):
        # train Line model with order
        self.alpha = self.init_alpha
        batch_size = self.batch_size
        t0 = time.time()
        num_batch = int(self.num_sampling_edge / batch_size)
        epoch_iter = tqdm(range(num_batch))
        for b in epoch_iter:
            if b % 100 == 0:
                epoch_iter.set_description(
                    f"Progress: {b *1.0/num_batch * 100:.4f}%, alpha: {self.alpha:.6f}, time: {time.time() - t0:.4f}"
                )
                self.alpha = self.init_alpha * max((1 - b * 1.0 / num_batch), 0.0001)

            for k in range(self.num_graph):
                u, v = [0] * batch_size, [0] * batch_size
                for i in range(batch_size):
                    edge_id = alias_draw(self.edges_prob[k][0], self.edges_prob[k][1])
                    u[i], v[i] = self.edges[k][edge_id]

                vec_error = np.zeros((batch_size, self.dimension))
                label, target = np.asarray([1 for i in range(batch_size)]), np.asarray(v)
                for j in range(1 + self.negative):
                    if j != 0:
                        label = np.asarray([0 for i in range(batch_size)])
                        for i in range(batch_size):
                            neg_node = alias_draw(self.node_prob[k][0], self.node_prob[k][1])
                            target[i] = self.id2node[k][neg_node]
                        self._update(self.emb_vertex[u], self.emb_context[target], vec_error, label)
                self.emb_vertex[u] += vec_error


#!/bin/env python

import ROOT as r
import numpy as np
from nesteddict import nesteddict as ndict
import json

antypes=[["E","e","Ele"],["Mu","mu","Mu"]]
form="output_CITo2%s_M%d_CUETP8M1_Lam%sTeV%s%s_13TeV_Pythia8_Corrected-v4_ntuple.root"
dyform="output_DYTo2%s_M%d_CUETP8M1_13TeV_Pythia8_Corrected-v3_ntuple.root"

r.gROOT.SetBatch(True)

mvals=[300,800,1300]
lvals=["1", "10", "16", "22", "28", "34", "100k"]
helis=["LL","LR","RR"]
intfs=["Con","Des"]

supers = [400,500,700,1100,1900,3500]

for antype in antypes:
    base="root://cmseos.fnal.gov///store/user/sturdy/ZprimeAnalysis/histosDec15/histosZprime{0:s}{0:s}/".format(antype[2])
    params = ndict()

    with open("ciparametrization_2%s.json"%(antype[1]),"w") as js:
        with open("cicounts_2%s.txt"%(antype[1]),"w") as out:
            for intf in intfs:
                for heli in helis:
                    files=[]
                    for point in supers[:-1]:
                        # params["%s%s_%dGeV"%(intf,heli,point)] = np.zeros(len(lvals),'float64')
                        params["%s%s_%dGeV"%(intf,heli,point)] = [0 for j in range(len(lvals))]
                    print("{0:s}{1:s}".format(intf,heli))
                    for i,lval in enumerate(lvals):
                        hist = None
                        can   = r.TCanvas("can","",800,800)
                        stack = r.THStack("stack","")
                        for mval in mvals:
                            lf = r.TFile.Open("%s%s"%(base,form%(antype[0],mval,lval,intf,heli)))
                            files.append(lf)  ## ROOT why are you so awful!!!!???
                            if not lf:
                                print(lf,hist)
                                continue
                            htmp = lf.Get("ZprimeRecomass")
                            htmp.Scale(1.3) # apply k-factor on signal samples
                            htmp = htmp.Rebin(100,"%s_%d_rebinned"%(htmp.GetName(),mval))
                            if mval == 300:
                                hist = htmp.Clone("htmp_%s%s%s%d"%(lval,intf,heli,mval))
                                htmp.SetLineColor(r.kOrange)
                                htmp.SetFillColor(r.kOrange)
                                htmp.SetFillStyle(3001)
                                stack.Add(htmp)
                            else:
                                hist.Add(htmp.Clone("htmp_%s%s%s%d"%(lval,intf,heli,mval)))
                                if mval == 1300:
                                    htmp.SetLineColor(r.kGreen)
                                    htmp.SetFillColor(r.kGreen)
                                    htmp.SetFillStyle(3001)
                                    stack.Add(htmp)
                                else:
                                    htmp.SetLineColor(r.kBlue)
                                    htmp.SetFillColor(r.kBlue)
                                    htmp.SetFillStyle(3001)
                                    stack.Add(htmp)
                                    pass
                                if mval == 1300 and "%s%s"%(intf,heli) == "ConLL":
                                    lf = r.TFile.Open("%s%s"%(base,form%(antype[0],2000,lval,intf,heli)))
                                    if not lf:
                                        print(lf,htmp,hist)
                                        continue
                                    htmp = lf.Get("ZprimeRecomass")
                                    htmp.Scale(1.3) # apply k-factor on signal samples
                                    htmp = htmp.Rebin(100,"%s_%d_rebinned"%(htmp.GetName(),2000))
                                    hist.Add(htmp.Clone("htmp_%s%s%s%d"%(lval,intf,heli,2000)))
                                    htmp.SetLineColor(r.kRed)
                                    htmp.SetFillColor(r.kRed)
                                    htmp.SetFillStyle(3001)
                                    stack.Add(htmp)
                                    pass
                                pass
                            pass
                        # stack.Rebin(10)
                        # hist.Rebin(10)
                        stack.SetMinimum(0.8*hist.GetMinimum(0.001))
                        stack.SetMaximum(1.25*hist.GetMaximum())
                        stack.Draw("hist")
                        stack.GetXaxis().SetRangeUser(0,5000)
                        stack.GetXaxis().SetNdivisions(505)
                        r.gPad.SetLogy(True)
                        hist.Draw("same")
                        can.Update()
                        # raw_input()
                        for ftype in ["png","C","pdf","eps"]:
                            can.SaveAs("~/public/forCIAnalysis/cito2{1:s}_{2:s}{3:s}{4:s}.{0:s}".format(ftype,antype[1],lval,intf,heli))
                        # raw_input("enter to continue")
                        for p,point in enumerate(supers[:-1]):
                            bval  = hist.FindBin(point)
                            upval = hist.FindBin(supers[p+1]-0.05)
                            val  = hist.Integral(bval,upval)
                            print("{0:s} {1:d} {2:d} {3:d} {4:2.4f}".format(lval,point,bval,upval,val))
                            out.write("{0:s} {1:d} {2:d} {3:d} {4:2.4f}\n".format(lval,point,bval,upval,val))
                            params["%s%s_%dGeV"%(intf,heli,point)][i] = val
                            pass
                        pass
                    pass
                pass
            print(params)
            pass
        json.dump(params,js)
        pass
    pass


from pepper.framework import AbstractMicrophone

from webrtcvad import Vad
import numpy as np

from Queue import Queue

from typing import Iterable

class Utterance(object):
    def __init__(self):
        """Create Utterance Object"""
        self._queue = Queue()
        self._frames = []

    @property
    def frames(self):
        # type: () -> Iterable[np.ndarray]
        """
        Get Utterance Frames

        Yields
        -------
        frames: Iterable of np.ndarray
        """

        if self._frames:
            for frame in self._frames:
                yield frame
        else:
            while True:
                frame = self._queue.get()
                if frame is None:
                    break
                self._frames.append(frame)
                yield frame

    @property
    def audio(self):
        # type: () -> np.ndarray
        """
        Get Utterance Audio (Concatenated Frames)

        Returns
        -------
        audio: np.ndarray
        """

        return np.concatenate([frame for frame in self.frames])

    def add_frame(self, frame):
        # type: (np.ndarray) -> None
        """
        Add Utterance Frame (done by VAD)

        Parameters
        ----------
        frame: np.ndarray
        """

        self._queue.put(frame)

    def __iter__(self):
        # type: () -> Iterable[np.ndarray]
        return self.frames

class VAD(object):

    AUDIO_FRAME_MS = 10
    BUFFER_SIZE = 100

    AUDIO_TYPE = np.int16
    AUDIO_TYPE_BYTES = 2

    VOICE_THRESHOLD = 0.6
    VOICE_WINDOW = 50

    MODE = 3

    def __init__(self, microphone):
        # type: (AbstractMicrophone) -> VAD
        """
        Perform Voice Activity Detection on Microphone Input

        Parameters
        ----------
        microphone: AbstractMicrophone
        """

        self._microphone = microphone

        self._vad = Vad(VAD.MODE)

        # Voice Activity Detection Frame Size, Atomic VAD Unit
        self._frame_size = VAD.AUDIO_FRAME_MS * self.microphone.rate // 1000
        self._frame_size_bytes = self._frame_size * VAD.AUDIO_TYPE_BYTES

        self._audio_buffer = np.zeros((VAD.BUFFER_SIZE, self._frame_size), VAD.AUDIO_TYPE)
        self._voice_buffer = np.zeros(VAD.BUFFER_SIZE, np.bool)
        self._buffer_index = 0

        self._utterance = None
        self._utterance_queue = Queue()

        self._frame_buffer = bytearray()

        self._activation = 0

        self.microphone.callbacks += [self._on_audio]

    @property
    def microphone(self):
        """
        VAD Microphone

        Returns
        -------
        microphone: AbstractMicrophone
        """
        return self._microphone

    @property
    def activation(self):
        """
        VAD Activation

        Returns
        -------
        activation: float
        """
        return self._activation

    @property
    def utterances(self):
        # type: () -> Iterable[Utterance]
        """
        Get Utterances from Microphone Stream

        Yields
        -------
        voices: Iterable of Voice
        """
        while True:
            yield self._utterance_queue.get()

    def _on_audio(self, audio):
        # type: (np.ndarray) -> None

        # Work through Microphone Stream Frame by Frame
        self._frame_buffer.extend(audio.tobytes())
        while len(self._frame_buffer) >= self._frame_size_bytes:
            self._on_frame(np.frombuffer(self._frame_buffer[:self._frame_size_bytes], VAD.AUDIO_TYPE))
            del self._frame_buffer[:self._frame_size_bytes]

    def _on_frame(self, frame):
        self._activation = self._calculate_activation(frame)

        if not self._utterance:
            if self.activation > VAD.VOICE_THRESHOLD:

                # Create New Utterance Object
                self._utterance = Utterance()

                # Add Buffer Contents to Utterance
                self._utterance.add_frame(self._audio_buffer[self._buffer_index:].ravel())
                self._utterance.add_frame(self._audio_buffer[:self._buffer_index].ravel())

                # Add Utterance to Utterance Queue
                self._utterance_queue.put(self._utterance)
        else:
            # If Utterance Ongoing: Add Frame to Utterance Object
            if self.activation > VAD.VOICE_THRESHOLD:
                self._utterance.add_frame(frame)

            # Else: Terminate Utterance
            else:
                self._utterance.add_frame(None)
                self._utterance = None

    def _calculate_activation(self, frame):
        # Update Buffers
        self._audio_buffer[self._buffer_index] = frame
        self._voice_buffer[self._buffer_index] = self._vad.is_speech(frame.tobytes(), self.microphone.rate, len(frame))
        self._buffer_index = (self._buffer_index + 1) % VAD.BUFFER_SIZE

        # Calculate Activation
        voice_window = np.arange(self._buffer_index - VAD.VOICE_WINDOW, self._buffer_index) % VAD.BUFFER_SIZE
        return np.mean(self._voice_buffer[voice_window])

    def __iter__(self):
        return self.utterances


from __future__ import division, print_function, absolute_import

import numpy as np
import numpy.testing as npt
import nibabel as nib
from nibabel.tmpdirs import InTemporaryDirectory
from dipy.io.streamline import (save_trk, load_trk, save_tractogram,
                                save_tck, load_tck, load_tractogram,
                                save_dpy, load_dpy)
from dipy.io.trackvis import save_trk as trackvis_save_trk
from dipy.tracking.streamline import Streamlines

streamline = np.array([[82.20181274,  91.36505891,  43.15737152],
                       [82.38442231,  91.79336548,  43.87036514],
                       [82.48710632,  92.27861023,  44.56298065],
                       [82.53310394,  92.78545381,  45.24635315],
                       [82.53793335,  93.26902008,  45.94785309],
                       [82.48797607,  93.75003815,  46.64939880],
                       [82.35533142,  94.25181581,  47.32533264],
                       [82.15484619,  94.76634216,  47.97451019],
                       [81.90982819,  95.28792572,  48.60244371],
                       [81.63336945,  95.78153229,  49.23971176],
                       [81.35479736,  96.24868011,  49.89558792],
                       [81.08713531,  96.69807434,  50.56812668],
                       [80.81504822,  97.14285278,  51.24193192],
                       [80.52591705,  97.56719971,  51.92168427],
                       [80.26599884,  97.98269653,  52.61848068],
                       [80.04635621,  98.38131714,  53.33855821],
                       [79.84691621,  98.77052307,  54.06955338],
                       [79.57667542,  99.13599396,  54.78985596],
                       [79.23351288,  99.43207551,  55.51065063],
                       [78.84815979,  99.64141846,  56.24016571],
                       [78.47383881,  99.77347565,  56.99299241],
                       [78.12837219,  99.81330872,  57.76969528],
                       [77.80438995,  99.85082245,  58.55574799],
                       [77.49439240,  99.88065338,  59.34777069],
                       [77.21414185,  99.85343933,  60.15090561],
                       [76.96416473,  99.82772827,  60.96406937],
                       [76.74712372,  99.80519104,  61.78676605],
                       [76.52263641,  99.79122162,  62.60765076],
                       [76.03757477, 100.08692169,  63.24152374],
                       [75.44867706, 100.35265351,  63.79513168],
                       [74.78033447, 100.57255554,  64.27278901],
                       [74.11605835, 100.77330781,  64.76428986],
                       [73.51222992, 100.98779297,  65.32373047],
                       [72.97387695, 101.23387146,  65.93502045],
                       [72.47355652, 101.49151611,  66.57343292],
                       [71.99834442, 101.72480774,  67.23979950],
                       [71.56909181, 101.98665619,  67.92664337],
                       [71.18083191, 102.29483795,  68.61888123],
                       [70.81879425, 102.63343048,  69.31127167],
                       [70.47422791, 102.98672485,  70.00532532],
                       [70.10092926, 103.28502655,  70.70999908],
                       [69.69512177, 103.51667023,  71.42147064],
                       [69.27423096, 103.71351624,  72.13452911],
                       [68.91260529, 103.81676483,  72.89796448],
                       [68.60788727, 103.81982422,  73.69258118],
                       [68.34162903, 103.76619720,  74.49915314],
                       [68.08542633, 103.70635223,  75.30856323],
                       [67.83590698, 103.60187531,  76.11553955],
                       [67.56822968, 103.44821930,  76.90870667],
                       [67.28399658, 103.25878906,  77.68825531],
                       [67.00117493, 103.03740692,  78.45989227],
                       [66.72718048, 102.80329895,  79.23099518],
                       [66.46197511, 102.54130554,  79.99622345],
                       [66.20803833, 102.22305298,  80.74387360],
                       [65.96872711, 101.88980865,  81.48987579],
                       [65.72864532, 101.59316254,  82.25085449],
                       [65.47808075, 101.33383942,  83.02194214],
                       [65.21841431, 101.11295319,  83.80186462],
                       [64.95678711, 100.94080353,  84.59326935],
                       [64.71759033, 100.82022095,  85.40114594],
                       [64.48053741, 100.73490143,  86.21411896],
                       [64.24304199, 100.65074158,  87.02709198],
                       [64.01773834, 100.55318451,  87.84204865],
                       [63.83801651, 100.41996765,  88.66333008],
                       [63.70982361, 100.25119019,  89.48779297],
                       [63.60707855, 100.06730652,  90.31262207],
                       [63.46164322,  99.91001892,  91.13648224],
                       [63.26287842,  99.78648376,  91.95485687],
                       [63.03713226,  99.68377686,  92.76905823],
                       [62.81192398,  99.56619263,  93.58140564],
                       [62.57145309,  99.42708588,  94.38592529],
                       [62.32259369,  99.25592804,  95.18167114],
                       [62.07497787,  99.05770111,  95.97154236],
                       [61.82253647,  98.83877563,  96.75438690],
                       [61.59536743,  98.59293365,  97.53706360],
                       [61.46530151,  98.30503845,  98.32772827],
                       [61.39904785,  97.97928619,  99.11172485],
                       [61.33279419,  97.65353394,  99.89572906],
                       [61.26067352,  97.30914307, 100.67123413],
                       [61.19459534,  96.96743011, 101.44847107],
                       [61.19580461,  96.63417053, 102.23215485],
                       [61.26572037,  96.29887391, 103.01185608],
                       [61.39840698,  95.96297455, 103.78307343],
                       [61.57207871,  95.64262391, 104.55268097],
                       [61.78163528,  95.35540771, 105.32629395],
                       [62.06700134,  95.09746552, 106.08564758],
                       [62.39427185,  94.85724641, 106.83369446],
                       [62.74076462,  94.62278748, 107.57482147],
                       [63.11461639,  94.40107727, 108.30641937],
                       [63.53397751,  94.20418549, 109.02002716],
                       [64.00019836,  94.03809357, 109.71183777],
                       [64.43580627,  93.87523651, 110.42416382],
                       [64.84857941,  93.69993591, 111.14715576],
                       [65.26740265,  93.51858521, 111.86515808],
                       [65.69511414,  93.36718751, 112.58474731],
                       [66.10470581,  93.22719574, 113.31711578],
                       [66.45891571,  93.06028748, 114.07256317],
                       [66.78582001,  92.90560913, 114.84281921],
                       [67.11138916,  92.79004669, 115.62040711],
                       [67.44729614,  92.75711823, 116.40135193],
                       [67.75688171,  92.98265076, 117.16111755],
                       [68.02041626,  93.28012848, 117.91371155],
                       [68.25725555,  93.53466797, 118.69052124],
                       [68.46047974,  93.63263702, 119.51107788],
                       [68.62039948,  93.62007141, 120.34690094],
                       [68.76782227,  93.56475067, 121.18331909],
                       [68.90222168,  93.46326447, 122.01765442],
                       [68.99872589,  93.30039978, 122.84759521],
                       [69.04119873,  93.05428314, 123.66156769],
                       [69.05086517,  92.74394989, 124.45450592],
                       [69.02742004,  92.40427399, 125.23509979],
                       [68.95466614,  92.09059143, 126.02339935],
                       [68.84975433,  91.79674531, 126.81564331],
                       [68.72673798,  91.53726196, 127.61715698],
                       [68.60685731,  91.30300141, 128.42681885],
                       [68.50636292,  91.12481689, 129.25317383],
                       [68.39311218,  91.01572418, 130.08976746],
                       [68.25946808,  90.94654083, 130.92756653]],
                      dtype=np.float32)

streamlines = Streamlines([streamline[[0, 10]], streamline,
                           streamline[::2], streamline[::3],
                           streamline[::5], streamline[::6]])

def test_io_streamline():
    with InTemporaryDirectory():
        fname = 'test.trk'
        affine = np.eye(4)

        # Test save
        save_tractogram(fname, streamlines, affine,
                        vox_size=np.array([2, 1.5, 1.5]),
                        shape=np.array([50, 50, 50]))
        tfile = nib.streamlines.load(fname)
        npt.assert_array_equal(affine, tfile.affine)
        npt.assert_array_equal(np.array([2, 1.5, 1.5]),
                               tfile.header.get('voxel_sizes'))
        npt.assert_array_equal(np.array([50, 50, 50]),
                               tfile.header.get('dimensions'))
        npt.assert_equal(len(tfile.streamlines), len(streamlines))
        npt.assert_array_almost_equal(tfile.streamlines[1], streamline,
                                      decimal=4)

        # Test basic save
        save_tractogram(fname, streamlines, affine)
        tfile = nib.streamlines.load(fname)
        npt.assert_array_equal(affine, tfile.affine)
        npt.assert_equal(len(tfile.streamlines), len(streamlines))
        npt.assert_array_almost_equal(tfile.streamlines[1], streamline,
                                      decimal=5)

        # Test Load
        local_streamlines, hdr = load_tractogram(fname)
        npt.assert_equal(len(local_streamlines), len(streamlines))
        for arr1, arr2 in zip(local_streamlines, streamlines):
            npt.assert_allclose(arr1, arr2)

def io_tractogram(load_fn, save_fn, extension):
    with InTemporaryDirectory():
        fname = 'test.{}'.format(extension)
        affine = np.eye(4)

        # Test save
        save_fn(fname, streamlines, affine, vox_size=np.array([2, 1.5, 1.5]),
                shape=np.array([50, 50, 50]))
        tfile = nib.streamlines.load(fname)
        npt.assert_array_equal(affine, tfile.affine)
        vox_size = tfile.header.get('voxel_sizes')
        dims = tfile.header.get('dimensions')
        if isinstance(vox_size, str):
                vox_size = vox_size.replace('[', '').replace(']', '')
                vox_size = np.fromstring(vox_size, sep=" ", dtype=np.float)
        if isinstance(dims, str):
                dims = dims.replace('[', '').replace(']', '')
                dims = np.fromstring(dims, sep=" ", dtype=np.int)
        npt.assert_array_equal(np.array([2, 1.5, 1.5]), vox_size)
        npt.assert_array_equal(np.array([50, 50, 50]), dims)
        npt.assert_equal(len(tfile.streamlines), len(streamlines))
        npt.assert_array_almost_equal(tfile.streamlines[1], streamline,
                                      decimal=4)

        # Test basic save
        save_fn(fname, streamlines, affine)
        tfile = nib.streamlines.load(fname)
        npt.assert_array_equal(affine, tfile.affine)
        npt.assert_equal(len(tfile.streamlines), len(streamlines))
        npt.assert_array_almost_equal(tfile.streamlines[1], streamline,
                                      decimal=5)

        # Test lazy save
        save_fn(fname, streamlines, affine, vox_size=np.array([2, 1.5, 1.5]),
                shape=np.array([50, 50, 50]), reduce_memory_usage=True)
        tfile = nib.streamlines.load(fname)
        npt.assert_array_equal(affine, tfile.affine)
        npt.assert_equal(len(tfile.streamlines), len(streamlines))
        npt.assert_array_almost_equal(tfile.streamlines[1], streamline,
                                      decimal=4)

        # Test Load
        local_streamlines, hdr = load_fn(fname)
        npt.assert_equal(len(local_streamlines), len(streamlines))
        for arr1, arr2 in zip(local_streamlines, streamlines):
            npt.assert_allclose(arr1, arr2, rtol=1e4)

        # Test lazy Load
        local_streamlines, hdr = load_fn(fname, lazy_load=True)
        for arr1, arr2 in zip(local_streamlines, streamlines):
            npt.assert_allclose(arr1, arr2, rtol=1e4)

def test_io_trk():
    io_tractogram(load_trk, save_trk, "trk")

def test_io_tck():
    io_tractogram(load_tck, save_tck, "tck")

def test_io_dpy():
    with InTemporaryDirectory():
        fname = 'test.dpy'

        # Test save
        save_dpy(fname, streamlines)
        tracks, _ = load_dpy(fname)
        npt.assert_equal(len(tracks), len(streamlines))
        npt.assert_array_almost_equal(tracks[1], streamline,
                                      decimal=4)

def test_trackvis():
    with InTemporaryDirectory():
        fname = 'trackvis_test.trk'
        affine = np.eye(4)

        # Test save
        trackvis_save_trk(fname, streamlines, affine, np.array([50, 50, 50]))
        tfile = nib.streamlines.load(fname)
        npt.assert_array_equal(affine, tfile.affine)
        npt.assert_array_equal(np.array([1., 1., 1.]),
                               tfile.header.get('voxel_sizes'))
        npt.assert_array_equal(np.array([50, 50, 50]),
                               tfile.header.get('dimensions'))
        npt.assert_equal(len(tfile.streamlines), len(streamlines))
        npt.assert_array_almost_equal(tfile.streamlines[1], streamline,
                                      decimal=4)

        # Test Deprecations
        npt.assert_warns(DeprecationWarning, trackvis_save_trk, fname,
                         streamlines, affine, np.array([50, 50, 50]))

if __name__ == '__main__':
    npt.run_module_suite()


import time
import math

import numpy
import ray
import torch

import models

@ray.remote
class SelfPlay:
    """
    Class which runs in a dedicated thread to play games and save them to the replay buffer.
    """

    def __init__(self, initial_checkpoint, Game, config, seed):
        self.config = config
        self.game = Game(seed)

        numpy.random.seed(seed)
        torch.manual_seed(seed)

        # Initialize MuZero's neural nets
        self.model = models.MuZeroNetwork(self.config)
        self.model.set_weights(initial_checkpoint["weights"])
        self.model.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
        self.model.eval()

        ### And that's it, pretty simple init.
        ### I tried running it again, but this time I got this error:
        ###     self.replay_buffer = replay_buffer.Reanalyse.options(
        ### AttributeError: module 'replay_buffer' has no attribute 'Reanalyse'
        ### So I went back to replay_buffer.py to implement the init for Reanalyze. By the way, he spelled reanalyze wrong lol.

    # This function is where the self play workers continuously play the game, saving the played games to the replay buffer for training.
    def continuous_self_play(self, shared_storage, replay_buffer, test_mode=False):
        while (
            ray.get(shared_storage.get_info.remote("training_step")) < self.config.training_steps
            and not ray.get(shared_storage.get_info.remote("terminate"))
        ):
            self.model.set_weights(ray.get(shared_storage.get_info.remote("weights")))

            # If not testing, play the game and save it to the replay buffer, otherwise take the best action (no exploration) in test mode
            if not test_mode:
                game_history = self.play_game(  # Play the game
                    # Although I think this multi-line argument is ugly, I also think it would be even uglier on one line, unfortunately.
                    # This softmax voodoo magic basically alters the visit count distribution to ensure that the action selection becomes greedier as training progresses.
                    # The smaller it is, the more likely the best action (ie with the highest visit count) is chosen.
                    self.config.visit_softmax_temperature_fn(
                        trained_steps=ray.get(
                            shared_storage.get_info.remote("training_step")
                        )
                    ),
                    self.config.temperature_threshold,
                    False,
                    "self",
                    0
                )
                replay_buffer.save_game.remote(game_history, shared_storage)  # Save the game to the replay buffer
            else:
                # Take the best action (no exploration) in test mode
                game_history = self.play_game(
                    0,
                    self.config.temperature_threshold,
                    False,
                    "self" if len(self.config.players) == 1 else self.config.opponent,
                    self.config.muzero_player  # This is a number indicating the turn Muzero begins to play (0: MuZero plays first, 1: MuZero plays second).
                )
                # Save game/episode to shared storage
                shared_storage.set_info.remote(
                    {
                        "episode_length": len(game_history.action_history) - 1,
                        "total_reward": sum(game_history.reward_history),
                        "mean_value": numpy.mean([value for value in game_history.root_values if value])
                    }
                )
                # If there's more than 1 players, save both MuZero's reward and the opponents reward, presumably to compare the two.
                if 1 < len(self.config.players):
                    shared_storage.set_info.remote(
                        {
                            "muzero_reward": sum(
                                # Sum rewards from reward history where the player is muzero.
                                reward for i, reward in enumerate(game_history.reward_history)
                                if game_history.to_play_history[i - 1] == self.config.muzero_player
                            ),
                            "opponent_reward": sum(
                                # Sum rewards from reward history where the player is NOT muzero.
                                # But wait, what if the game is a 3-player game? Will ALL of muzero's opponents' rewards be summed together? That seems unfair to muzero.
                                reward for i, reward in enumerate(game_history.reward_history)
                                if game_history.to_play_history[i - 1] != self.config.muzero_player
                            )
                        }
                    )

            # Similar to what we saw in trainer.py, here we adjust the self play / training ratio to avoid over/underfitting.
            if not test_mode and self.config.self_play_delay:
                time.sleep(self.config.self_play_delay)
            if not test_mode and self.config.ratio:
                while (  # Another really large condition for a while loop that just puts the program to sleep for a bit.
                    ray.get(shared_storage.get_info.remote("training_step")) / max(1, ray.get(shared_storage.get_info.remote("num_played_steps"))) < self.config.ratio
                    and ray.get(shared_storage.get_info.remote("training_step")) < self.config.training_steps
                    and not ray.get(shared_storage.get_info.remote("terminate"))
                ):
                    time.sleep(0.5)

        self.close_game()

    def play_game(self, temperature, temperature_threshold, render, opponent, muzero_player):
        """
        Play one game with actions based on Monte Carlo Tree Search at each move.

        :param temperature:             Used in selecting the action. From the paper, "visit count distribution is parametrized using a temperature" (Appendix D Data Generation)
        :param temperature_threshold:   Also used in selecting the action, I think its the max value that the temperature can be?
        :param render:                  Boolean, whether or not to render the game
        :param opponent:                The opponent to play against, either "self" or self.config.opponent
        :param muzero_player:           The trained muzero model to play against. FIXME: If muzero_player is 0 and opponent is not self, what happens?
        :return:                        game_history
        """

        game_history = GameHistory()
        observation = self.game.reset()
        game_history.action_history.append(0)  # Append action_history with 0? Why?
        ### Since this thing depends on the GameHistory class, I went down and implemented that. It's not too long.
        game_history.observation_history.append(observation)
        game_history.reward_history.append(0)
        game_history.to_play_history.append(self.game.to_play())  # to_play() determines the current player. It seems to be defined in both the game file and the Node class, but sometimes not in the game file? cartpole doesn't define to_play, so what happens there? is this just for multiplayer games?

        done = False

        # Render initial state of game/environment.
        if render:
            self.game.render()

        with torch.no_grad():  # with statement in Python is used in exception handling. Not sure what torch.no_grad is used for in this though.
            # Main gameplay loop.
            while not done and len(game_history.action_history) <= self.config.max_moves:
                # Make sure observation is 3 dimensional? Does it have to be 3 dimensional? What if the observation is a string, a one-dimensional array of characters?
                # Ohhh, I think the observation has to be 3 dimensional because the resnet is a 3d convolutional network? Can we change it to have a one dimensional input, though?
                assert (len(numpy.array(observation).shape) == 3), f"Observation should be 3 dimensional instead of {len(numpy.array(observation).shape)} dimensional. Got observation of shape: {numpy.array(observation).shape}"
                # Make sure the observation shape matches the observation_shape defined in MuZeroConfig. Makes sense.
                assert (numpy.array(observation).shape == self.config.observation_shape), f"Observation should match the observation_shape defined in MuZeroConfig. Expected {self.config.observation_shape} but got {numpy.array(observation).shape}."
                stacked_observations = game_history.get_stacked_observations(-1, self.config.stacked_observations)  # config.stacked_observations is the number of previous observations and previous actions to add to the current observation

                # Choose the action
                # If it's playing against itself, use select_action(), if it's playing against a preprogrammed or human opponent, use select_opponent_action().
                if opponent == "self" or muzero_player == self.game.to_play():
                    root, mcts_info = MCTS(self.config).run(self.model, stacked_observations, self.game.legal_actions(), self.game.to_play(), True)  # Create and run the MCTS.
                    # We input the root node because select_action() gets the visit counts of all the root node's children.
                    action = self.select_action(root, temperature if not temperature_threshold or len(game_history.action_history) < temperature_threshold else 0)
                    # If we're rendering the game, also print out information about the MCTS.
                    if render:
                        print(f'Tree depth: {mcts_info["max_tree_depth"]}')
                        print(f"Root value for player {self.game.to_play()}: {root.value():.2f}")
                else:
                    action, root = self.select_opponent_action(opponent, stacked_observations)

                # After choosing the action, execute it in the environment and from it get the observation, reward, and whether or not we're done.
                observation, reward, done = self.game.step(action)

                # Print the action taken and render game state after action has been taken.
                if render:
                    print(f"Played action: {self.game.action_to_string(action)}")
                    self.game.render()

                # Store tree search statistics of game to game_history.
                game_history.store_search_statistics(root, self.config.action_space)

                # Store next batch
                game_history.action_history.append(action)
                game_history.observation_history.append(observation)
                game_history.reward_history.append(reward)
                game_history.to_play_history.append(self.game.to_play())

        # Return game_history, which is saved in the replay buffer and shared storage. FIXME: Shared storage saves to disk, but what about the replay buffer? Is that our bottleneck?
        return game_history

    # Is this function really necessary? Why not just use self.game.close() in continuous_self_play()? It's only one line.
    def close_game(self):
        self.game.close()

    # Select the action for the opponent playing against muzero.
    def select_opponent_action(self, opponent, stacked_observations):
        """
        Select the action for the opponent playing against muzero.

        :param opponent:                The opponent playing against muzero.
        :param stacked_observations:
        :return:                        action
        """
        # Select action depending on the type of opponent.
        if opponent == "human":
            # Create a MuZero MCTS so that MuZero can suggest a move to the human player.
            root, mcts_info = MCTS(self.config).run(
                self.model,
                stacked_observations,
                self.game.legal_actions(),
                self.game.to_play(),
                True,
            )
            # Print out info about the MCTS, let MuZero suggest its move to the human player.
            print(f'Tree depth: {mcts_info["max_tree_depth"]}')
            print(f"Root value for player {self.game.to_play()}: {root.value():.2f}")
            print(f"Player {self.game.to_play()} turn. MuZero suggests {self.game.action_to_string(self.select_action(root, 0))}")
            return self.game.human_to_action(), root
        elif opponent == "expert":
            return self.game.expert_agent(), None
        elif opponent == "random":
            # Make sure legal_actions are valid.
            assert (self.game.legal_actions()), f"Legal actions should not be an empty array. Got {self.game.legal_actions()}."
            assert set(self.game.legal_actions()).issubset(set(self.config.action_space)), "Legal actions should be a subset of the action space."
            return numpy.random.choice(self.game.legal_actions()), None
        else:
            raise NotImplementedError(f'Wrong argument: "opponent" argument should be "self", "human", "expert" or "random". "{opponent}" is not an implemented opponent.')

        # Select MuZero's action based on the visit counts of each node in the MCTS. But where do the neural nets come into play?
    @staticmethod
    def select_action(node, temperature):
        """
        Select action according to the visit count distribution and the temperature.
        The temperature is changed dynamically with the visit_softmax_temperature function in the config.

        :param node:            We will get all the visit counts for all of this node's child nodes so we can calculate the visit_count_distribution to select the action.
        :param temperature:     Some value used in the voodoo we use to determine the visit_count_distribution
        :return:                action
        """
        # Get visit counts for all of the input node's child nodes.
        visit_counts = numpy.array([child.visit_count for child in node.children.values()], dtype="int32")
        # Make a list of actions which have resulted in a node. Is this a list of all possible actions?
        actions = [action for action in node.children.keys()]  # node.children is a dictionary, where the child nodes are the values and each child node's key is the action taken to get that child node.
        # If temperature == 0, select the action of the node with the most visits, if temperature == infinity, select the action randomly. If the temperature is somewhere in between, use some voodoo with the temperature and visit counts to select the action.
        if temperature == 0:
            action = actions[numpy.argmax(visit_counts)]
        elif temperature == float("inf"):
            action = numpy.random.choice(actions)
        else:
            # Some voodoo to get the distribution of the visit count of each action. See paper appendix Data Generation
            visit_count_distribution = visit_counts ** (1 / temperature)
            visit_count_distribution = visit_count_distribution / sum(visit_count_distribution)
            action = numpy.random.choice(actions, p=visit_count_distribution)  # The probability for each action is determined by the visit_count_distribution.

        return action

# This MCTS is game independent, so you can use it to play any game.
class MCTS:
    """
    Core Monte Carlo Tree Search algorithm.
    To decide on an action, we run N simulations, always starting at the root of the search tree and traversing the tree according to the UCB formula until we reach a
    leaf node. (UCB = Upper-Confidence-Bound, which is some voodoo used in MCTS, see https://en.wikipedia.org/wiki/Monte_Carlo_tree_search#Monte_Carlo_Method)
    """

    # Not much in the config, the real meat is in run()
    def __init__(self, config):
        self.config = config

    # Runs the MCTS, usually called at the same time the MCTS is initialized. to_play is the current player. override_root_with is the node to override the root node with.
    # FIXME: I still don't entirely understand this thing and exactly how it's used. You'll see my comments expounding my confusion below.
    def run(self, model, observation, legal_actions, to_play, add_exploration_noise, override_root_with=None):
        """
        At the root of the search tree we use the representation function to obtain a hidden state of the current observation.
        We then run a Monte Carlo Tree Search using only the action sequences and the model learned by the network. Cool, right?
        """
        if override_root_with:
            root = override_root_with
            root_predicted_value = None
        else:
            # Set up and expand the root node.
            root = Node(0)  # Initialize the root node
            observation = (
                torch.tensor(observation)
                .float()
                .unsqueeze(0)
                .to(next(model.parameters()).device)
            )
            root_predicted_value, reward, policy_logits, hidden_state = model.initial_inference(observation)  # model.initial_inference uses the representation function to encode MuZero's first (aka initial) observation.
            ### models.support_to_scalar is used in here, but I hadn't implemented that yet, so let's go over to models.py and implement it now!
            root_predicted_value = models.support_to_scalar(root_predicted_value, self.config.support_size).item()
            reward = models.support_to_scalar(reward, self.config.support_size).item()
            assert legal_actions, f"Legal actions should not be an empty array. Got {legal_actions}."  # Assert legal_actions exists.
            assert set(legal_actions).issubset(set(self.config.action_space)), "Legal actions should be a subset of the action space."
            root.expand(legal_actions, to_play, reward, policy_logits, hidden_state)

        if add_exploration_noise:
            root.add_exploration_noise(dirichlet_alpha=self.config.root_dirichlet_alpha, exploration_fraction=self.config.root_exploration_fraction)

        min_max_stats = MinMaxStats()  # So this is what holds the "min-max values of the tree", so I guess that means it holds the best actions the tree predicts you could take?

        max_tree_depth = 0  # Initialize the depth of the tree.
        # With the root set up, FIXME: run as many tree searches as we set for num_simulations? Really, we're running more than one MCTS? I'm not sure if that's awesome or unnecessary.
        # FIXME: Wait, I'm confused, are we running MULTIPLE tree searches here, or is this loop just used to expand a single tree?
        #  Based on what I know I would assume the latter, but why does the for loop set the node to root each time? Are we making the tree branch by branch rather than layer by layer?
        #  But what about branches that have the same starting nodes?
        for _ in range(self.config.num_simulations):
            virtual_to_play = to_play
            node = root
            search_path = [node]  # So is this to hold a branch of the tree? Or an entire tree?
            current_tree_depth = 0

            # Run the rest of the tree search. Not actually in this while loop, though, this while loop is just used to... what?
            # I think this is used to create a single branch? Not the rest of the tree search? But wouldn't that need the dynamics function?
            # Wait, or is this used to skip through the nodes that we've already created in a branch and skip to the leaf of the branch so we can expand it? Is that it?
            while node.expanded():
                current_tree_depth += 1  # Increment tree depth
                action, node = self.select_child(node, min_max_stats)  # Select best action and determine the next node from the current node based on the min_max_stats.
                search_path.append(node)

                # Players play turn by turn. But wait, wouldn't this if-else only work for two player games? What about 3+ players?
                # Oh, wait, below it says "More than two player mode not implemented." Ok. Whatever, our game is 1 player anyway
                if virtual_to_play + 1 < len(self.config.players):
                    virtual_to_play = self.config.players[virtual_to_play + 1]
                else:
                    virtual_to_play = self.config.players[0]

            # Here's the REALLY cool part:
            # Inside the search tree we use the dynamics function to obtain the next hidden state given an action and the previous hidden state.
            parent = search_path[-2]
            # model.recurrent_inference is where we call the dynamics function to obtain each hidden state from its previous state, including the hidden state from the representation function from the initial observation.
            value, reward, policy_logits, hidden_state = model.recurrent_inference(parent.hidden_state, torch.tensor([[action]]).to(parent.hidden_state.device))
            value = models.support_to_scalar(value, self.config.support_size).item()
            reward = models.support_to_scalar(reward, self.config.support_size).item()
            node.expand(self.config.action_space, virtual_to_play, reward, policy_logits, hidden_state)  # create a new node from the previous one?

            self.backpropagate(search_path, value, virtual_to_play, min_max_stats)  # Backpropagate to improve the network or tree or something.

            max_tree_depth = max(max_tree_depth, current_tree_depth)  # The depth of the biggest tree we've made here.

        extra_info = {
            "max_tree_depth": max_tree_depth,
            "root_predicted_value": root_predicted_value
        }
        return root, extra_info  # FIXME: So why are we returning the root? Are we... using the root value to determine the next move or something?

    def select_child(self, node, min_max_stats):
        """
        Select the child with the highest UCB score, aka the best action.
        """
        max_ucb = max(self.ucb_score(node, child, min_max_stats) for action, child in node.children.items())  # Find the action with the highest ucb score of the node's children.
        # Select an action randomly from a list of children with the highest ucb score. So wait, wouldn't that list only have one item? Unless there are two children with the max score.
        action = numpy.random.choice([action for action, child in node.children.items() if self.ucb_score(node, child, min_max_stats) == max_ucb])
        return action, node.children[action]

    def ucb_score(self, parent, child, min_max_stats):
        """
        Find the ucb score for a node based on its value, plus an exploration bonus based on the prior. I'm not sure what a prior is.
        """
        # Basically this is all voodoo, so unfortunately I don't have much to say about it.
        # The score seems to mostly be based on the node's visit count, so where do the neural nets come into play?
        # OH! The neural nets come into play in backpropagate(), as backpropagate() uses the reward and value from the dynamics function to improve the min_max_stats! Great!
        pb_c = (
            math.log(
                (parent.visit_count + self.config.pb_c_base + 1) / self.config.pb_c_base
            )
            + self.config.pb_c_init
        )
        pb_c *= math.sqrt(parent.visit_count) / (child.visit_count + 1)

        prior_score = pb_c * child.prior

        if child.visit_count > 0:
            # Mean value Q
            # Is this some kind of Q-learning thing?
            value_score = min_max_stats.normalize(
                child.reward
                + self.config.discount
                * (child.value() if len(self.config.players) == 1 else -child.value())
            )
        else:
            value_score = 0

        return prior_score + value_score

    def backpropagate(self, search_path, value, to_play, min_max_stats):
        """
        At the end of a simulation, we propagate the evaluation all the way up the tree to the root.
        So is a simulation an exploration of a tree branch? So would num_simulations be the number of branches in a tree?
        In that case, it makes perfect sense why you can still get such great results with MuZero with very few simulations, because the neural nets figure out which
        branches are the most promising before we even get started, so we don't really need that many, as the ones the neural nets come up with will already be really good.
        """
        if len(self.config.players) == 1:
            # For each node in the search path, update the min_max_stats with the reward and value from the dynamics function.
            for node in reversed(search_path):
                node.value_sum += value  # Add value from dynamics function to value_sum
                node.visit_count += 1  # Increment node visit count.
                min_max_stats.update(node.reward + self.config.discount * node.value())  # Update the min_max_stats with the reward and value from the dynamics function! So this is where the NNs come into play!
                # Adjust value with reward and discount values.
                value = node.reward + self.config.discount * value
        elif len(self.config.players) == 2:
            for node in reversed(search_path):
                node.value_sum += value if node.to_play == to_play else -value  # Add -value for opponent's value.
                node.visit_count += 1
                min_max_stats.update(node.reward + self.config.discount * -node.value())
                value = (-node.reward if node.to_play == to_play else node.reward) + self.config.discount * value
        else:
            raise NotImplementedError("More than two player mode not implemented.")

class Node:
    def __init__(self, prior):
        self.visit_count = 0
        self.to_play = -1
        self.prior = prior
        self.value_sum = 0
        self.children = {}
        self.hidden_state = None
        self.reward = 0

    # Return whether or not the node is expanded, i.e. it has children.
    def expanded(self):
        return len(self.children) > 0

    # Return the TRUE value of node, which is actually the value_sum / visit_count. So the less visited the node is, and the higher its value, means it's overall value is higher, which leads to muzero exploring promising, yet less tested nodes. Cool!
    def value(self):
        if self.visit_count == 0:
            return 0
        return self.value_sum / self.visit_count

    # Here's a lot of the meat of how the MCTS is created, and this is where the information from the neural nets is used the most.
    def expand(self, actions, to_play, reward, policy_logits, hidden_state):
        """
        Expand a node, i.e. create new children for it, using the value, reward, and policy prediction from the neural network! Cool, so this is also where the neural nets come into play!

        :param actions:         Legal actions
        :param to_play:         Player number
        :param reward:          The predicted reward from either the initial_inference (which uses some weird voodoo I don't understand) or the recurrent_inference using the dynamics function.
        :param policy_logits:   The policy/action to take.
        :param hidden_state:    The hidden state from either the representation function or the dynamics function.
        :return:                void, no return
        """
        self.to_play = to_play
        self.reward = reward
        self.hidden_state = hidden_state

        policy_values = torch.softmax(torch.tensor([policy_logits[0][a] for a in actions]), dim=0).tolist()  # Find values for each action using the policy,
        policy = {a: policy_values[i] for i, a in enumerate(actions)}  # use values to select best actions,
        for action, p in policy.items():
            self.children[action] = Node(p)  # and use make a node from each of the best actions! Awesome!

    # At the start of each search, add noise to the prior of the root to encourage the search to explore new actions.
    def add_exploration_noise(self, dirichlet_alpha, exploration_fraction):
        """
        At the start of each search, we add dirichlet noise to the prior of the root to encourage the search to explore new actions.
        The way we use dirichlet noise and why is explained well in this video: https://youtu.be/L0A86LmH7Yw?t=697
        """
        actions = list(self.children.keys())
        noise = numpy.random.dirichlet([dirichlet_alpha] * len(actions))
        frac = exploration_fraction
        for a, n in zip(actions, noise):
            self.children[a].prior = self.children[a].prior * (1 - frac) + n * frac

class GameHistory:
    """
    Store only useful information of a self-play game. FIXME: Where does it store it? RAM? I hope it doesn't store it there permanently...
    """

    def __init__(self):
        self.observation_history = []
        self.action_history = []
        self.reward_history = []
        self.to_play_history = []
        self.child_visits = []
        self.root_values = []
        self.reanalysed_predicted_root_values = None
        # For Prioritized Replay
        self.priorities = None
        self.game_priority = None

    # Store values of roots of each tree to root_values.
    def store_search_statistics(self, root, action_space):
        # Turn visit count from root into a policy, aka a set of best actions I think.
        if root is not None:
            sum_visits = sum(child.visit_count for child in root.children.values())
            # Append child_visits... with what and why?
            self.child_visits.append(
                [
                    root.children[a].visit_count / sum_visits
                    if a in root.children
                    else 0
                    for a in action_space
                ]
            )
            self.root_values.append(root.value())
        else:
            self.root_values.append(None)

    # Get the previous actions and observations to add to the current observation.
    def get_stacked_observations(self, index, num_stacked_observations):
        """
        Generate a new observation with the observation at the index position and num_stacked_observations past observations and actions stacked.

        :param index:                       Index from which we start stacking the observations.
        :param num_stacked_observations:    Number of past observations and actions stacked.
        :return:                            stacked_observations
        """
        # Convert index to positive value
        index = index % len(self.observation_history)

        stacked_observations = self.observation_history[index].copy()  # Copy observations from index in observation_history.
        # For all past observations from the index to the index - num_stacked_observations, get previous observations
        for past_observation_index in reversed(range(index - num_stacked_observations, index)):
            # Not exactly sure what this voodoo is for.
            if 0 <= past_observation_index:
                previous_observation = numpy.concatenate(  # numpy.concatenate concatenates a tuple of arrays, hence the seemingly redundant double parentheses.
                    (
                        self.observation_history[past_observation_index],
                        [numpy.ones_like(stacked_observations[0]) * self.action_history[past_observation_index + 1]]
                    )
                )
            else:
                previous_observation = numpy.concatenate(
                    (
                        numpy.zeros_like(self.observation_history[index]),
                        [numpy.zeros_like(stacked_observations[0])]
                    )
                )

            stacked_observations = numpy.concatenate((stacked_observations, previous_observation))

        return stacked_observations
            ### Alright, we'll figure out exactly what this stuff does later, for now, back to play_game()

# The word "stats" makes it sound like a big set of values or a probability distribution or something, but no, it's literally two float values: minimum and maximum.
class MinMaxStats:
    """
    A class that holds the min-max values of the tree. Literally two float values: minimum and maximum.
    """

    def __init__(self):
        self.maximum = -float("inf")
        self.minimum = float("inf")

    def update(self, value):
        self.maximum = max(self.maximum, value)
        self.minimum = min(self.minimum, value)

    def normalize(self, value):
        if self.maximum > self.minimum:
            # We normalize only when we have set the maximum and minimum values no maximum will be greater than the default infinity that the min is set to. Clever.
            return (value - self.minimum) / (self.maximum - self.minimum)
        return value

    ### Wow, despite all my comments, this file is STILL slightly shorter than the original because I didn't expand every parentheses in sight.
    ### Anyway, we're done with this whole thing! Awesome, let's head back over to muzero.py now!


"""
Implementation of Blackjack.
Note: Blackjacks are ignored since no decision is needed. Splits not implemented.
"""

import datetime
import os
import random

import numpy as np
import torch

from .abstract_game import AbstractGame

class MuZeroConfig:
    def __init__(self):

        self.seed = 0  # Seed for numpy, torch and the game
        self.max_num_gpus = 1  # Fix the maximum number of GPUs to use. It's usually faster to use a single GPU (set it to 1) if it has enough memory. None will use every GPUs available

        # Game
        self.observation_shape = (3, 3,
                                  3)  # Dimensions of the game observation, must be 3D (channel, height, width). For a 1D array, please reshape it to (1, 1, length of array)
        self.action_space = list(range(4))  # Fixed list of all possible actions. You should only edit the length
        self.players = list(range(1))  # List of players. You should only edit the length
        self.stacked_observations = 0  # Number of previous observations and previous actions to add to the current observation

        # Evaluate
        self.muzero_player = 0  # Turn Muzero begins to play (0: MuZero plays first, 1: MuZero plays second)
        self.opponent = None  # Hard coded agent that MuZero faces to assess his progress in multiplayer games. It doesn't influence training. None, "random" or "expert" if implemented in the Game class

        ### Self-Play
        self.num_workers = 4  # Number of simultaneous threads/workers self-playing to feed the replay buffer
        self.selfplay_on_gpu = False
        self.max_moves = 21  # Maximum number of moves if game is not finished before
        self.num_simulations = 21  # Number of future moves self-simulated
        self.discount = 1  # Chronological discount of the reward
        self.temperature_threshold = None  # Number of moves before dropping the temperature given by visit_softmax_temperature_fn to 0 (ie selecting the best action). If None, visit_softmax_temperature_fn is used every time

        # Root prior exploration noise
        self.root_dirichlet_alpha = 0.25
        self.root_exploration_fraction = 0.25

        # UCB formula
        self.pb_c_base = 19652
        self.pb_c_init = 1.25

        ### Network
        self.network = "resnet"  # "resnet" / "fullyconnected"
        self.support_size = 10  # Value and reward are scaled (with almost sqrt) and encoded on a vector with a range of -support_size to support_size. Choose it so that support_size <= sqrt(max(abs(discounted reward)))

        # Residual Network
        self.downsample = False  # Downsample observations before representation network, False / "CNN" (lighter) / "resnet" (See paper appendix Network Architecture)
        self.blocks = 2  # Number of blocks in the ResNet
        self.channels = 32  # Number of channels in the ResNet
        self.reduced_channels_reward = 32  # Number of channels in reward head
        self.reduced_channels_value = 32  # Number of channels in value head
        self.reduced_channels_policy = 32  # Number of channels in policy head
        self.resnet_fc_reward_layers = [16]  # Define the hidden layers in the reward head of the dynamic network
        self.resnet_fc_value_layers = [16]  # Define the hidden layers in the value head of the prediction network
        self.resnet_fc_policy_layers = [16]  # Define the hidden layers in the policy head of the prediction network

        # Fully Connected Network
        self.encoding_size = 32
        self.fc_representation_layers = [16]  # Define the hidden layers in the representation network
        self.fc_dynamics_layers = [16]  # Define the hidden layers in the dynamics network
        self.fc_reward_layers = [16]  # Define the hidden layers in the reward network
        self.fc_value_layers = [16]  # Define the hidden layers in the value network
        self.fc_policy_layers = [16]  # Define the hidden layers in the policy network

        ### Training
        self.results_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "../results",
                                         os.path.basename(__file__)[:-3], datetime.datetime.now().strftime(
                "%Y-%m-%d--%H-%M-%S"))  # Path to store the model weights and TensorBoard logs
        self.save_model = True  # Save the checkpoint in results_path as model.checkpoint
        self.training_steps = 50000  # Total number of training steps (ie weights update according to a batch)
        self.batch_size = 64  # Number of parts of games to train on at each training step
        self.checkpoint_interval = 10  # Number of training steps before using the model for self-playing
        self.value_loss_weight = 0.25  # Scale the value loss to avoid overfitting of the value function, paper recommends 0.25 (See paper appendix Reanalyze)
        self.train_on_gpu = torch.cuda.is_available()  # Train on GPU if available

        self.optimizer = "Adam"  # "Adam" or "SGD". Paper uses SGD
        self.weight_decay = 1e-4  # L2 weights regularization
        self.momentum = 0.9  # Used only if optimizer is SGD

        # Exponential learning rate schedule
        self.lr_init = 0.1  # Initial learning rate
        self.lr_decay_rate = 0.9  # Set it to 1 to use a constant learning rate
        self.lr_decay_steps = 2000

        ### Replay Buffer
        self.replay_buffer_size = 10000  # Number of self-play games to keep in the replay buffer
        self.num_unroll_steps = 20  # Number of game moves to keep for every batch element
        self.td_steps = 50  # Number of steps in the future to take into account for calculating the target value
        self.PER = True  # Prioritized Replay (See paper appendix Training), select in priority the elements in the replay buffer which are unexpected for the network
        self.PER_alpha = 1  # 0.5  # How much prioritization is used, 0 corresponding to the uniform case, paper suggests 1

        # Reanalyze (See paper appendix Reanalyse)
        self.use_last_model_value = True  # Use the last model to provide a fresher, stable n-step value (See paper appendix Reanalyze)
        self.reanalyse_on_gpu = False

        ### Adjust the self play / training ratio to avoid over/underfitting
        self.self_play_delay = 0  # Number of seconds to wait after each played game
        self.training_delay = 0  # Number of seconds to wait after each training step
        self.ratio = None  # Desired training steps per self played step ratio. Equivalent to a synchronous version, training can take much longer. Set it to None to disable it

    def visit_softmax_temperature_fn(self, trained_steps):
        """
        Parameter to alter the visit count distribution to ensure that the action selection becomes greedier as training progresses.
        The smaller it is, the more likely the best action (ie with the highest visit count) is chosen.
        Returns:
            Positive float.
        """
        if trained_steps < 500e3:
            return 1.0
        elif trained_steps < 750e3:
            return 0.5
        else:
            return 0.25

class Game(AbstractGame):
    """
    Game wrapper.
    """

    def __init__(self, seed=None):
        self.env = Blackjack(seed, 1)

    def step(self, action):
        """
        Apply action to the game.
        Args:
            action : action of the action_space to take.
        Returns:
            The new observation, the reward and a boolean if the game has ended.
        """
        observation, reward, done = self.env.step(action)
        return observation, reward, done

    def to_play(self):
        """
        Return the current player.
        Returns:
            The current player, it should be an element of the players list in the config.
        """
        return self.env.to_play()

    def legal_actions(self):
        """
        Should return the legal actions at each turn, if it is not available, it can return
        the whole action space. At each turn, the game have to be able to handle one of returned actions.
        For complex game where calculating legal moves is too long, the idea is to define the legal actions
        equal to the action space but to return a negative reward if the action is illegal.
        Returns:
            An array of integers, subset of the action space.
        """
        return self.env.legal_actions()

    def reset(self):
        """
        Reset the game for a new game.
        Returns:
            Initial observation of the game.
        """
        return self.env.reset()

    def render(self):
        """
        Display the game observation.
        """
        self.env.render()
        input("Press enter to take a step ")

    def human_to_action(self):
        """
        For multiplayer games, ask the user for a legal action
        and return the corresponding action number.
        Returns:
            An integer from the action space.
        """
        choice = input(
            f"Enter the action (0) Hit, (1) Stand, (2) Double down, or (3) Surrender "
            f"for the player {self.to_play()}: "
        )
        while choice not in [str(action) for action in self.legal_actions()]:
            print(f'{choice} is not a legal action. Choose another.')
            choice = input("Enter either (0) Hit, (1) Stand, (2) Double down, or (3) Surrender : ")
        return int(choice)

    def action_to_string(self, action_number):
        """
        Convert an action number to a string representing the action.
        Args:
            action_number: an integer from the action space.
        Returns:
            String representing the action.
        """
        actions = {
            0: "Hit",
            1: "Stand",
            2: "Double down",
            3: "Surrender"
        }
        return f"{action_number}. {actions[action_number]}"

card_values = {
    "2": 2,
    "3": 3,
    "4": 4,
    "5": 5,
    "6": 6,
    "7": 7,
    "8": 8,
    "9": 9,
    "10": 10,
    "J": 10,
    "Q": 10,
    "K": 10,
    "A": 11,
}

class Card:

    def __init__(self, face: str):
        self.face = face
        self.value = card_values[face]

    def __str__(self):
        return self.face

    def __repr__(self):
        return str(self.face)

class Blackjack:

    def __init__(self, seed: int, num_decks: int):
        self.num_decks = num_decks
        self.first_round = True
        self.player_round = True
        self.bet_multiplier = 1
        self.dealer_hand = []
        self.player_hand = []
        self.random = np.random.default_rng(seed=seed)
        self.deck = {
            "2": num_decks * 4,
            "3": num_decks * 4,
            "4": num_decks * 4,
            "5": num_decks * 4,
            "6": num_decks * 4,
            "7": num_decks * 4,
            "8": num_decks * 4,
            "9": num_decks * 4,
            "10": num_decks * 4,
            "J": num_decks * 4,
            "Q": num_decks * 4,
            "K": num_decks * 4,
            "A": num_decks * 4,
        }

        self.player_hand.append(self.deal_card())
        self.dealer_hand.append(self.deal_card())
        self.player_hand.append(self.deal_card())
        self.dealer_hand.append(self.deal_card())

        # if self.get_hand_value(self.player_hand) == 21:
        #     self.reset()

        self.player = 1

    def to_play(self):
        return 0 if self.player == 1 else 1

    def reset(self):
        self.deck = {k: 4 * self.num_decks for k, _ in self.deck.items()}
        self.player_hand = []
        self.dealer_hand = []
        self.player_hand.append(self.deal_card())
        self.dealer_hand.append(self.deal_card())
        self.player_hand.append(self.deal_card())
        self.dealer_hand.append(self.deal_card())
        self.player = 1
        # if self.get_hand_value(self.player_hand) == 21:
        #     self.reset()
        return self.get_observation()

    """
    Action: 0 = Hit
    Action: 1 = Stand
    Action 2 = Double down
    Action 3 = Surrender
    """

    def step(self, action):
        natural = False
        if self.get_hand_value(self.player_hand) == 21:
            natural = True

        if action == 0:
            self.player_hand.append(self.deal_card())

        if action == 2:
            self.bet_multiplier = 2
            self.player_hand.append(self.deal_card())

        done = self.is_busted() or action == 1 or action == 2 or action == 3 or self.get_hand_value(
            self.player_hand) == 21

        if done and not natural:
            self.player_round = False
            self.dealer_plays()

        if self.first_round and not natural:
            self.first_round = False

        return self.get_observation(), self.get_reward(action, done) * 100, done

    def get_observation(self):
        return np.array([
            np.full((3, 3), self.get_hand_value(self.player_hand), dtype="float32"),
            np.full((3, 3), self.dealer_hand[0].value, dtype="float32"),
            np.full((3, 3), 0),
        ])

    def legal_actions(self):
        # 0 = hit
        # 1 = stand
        # 2 = Double down (only first round)
        # 3 = Surrender (only first round)
        if self.first_round:
            return [0, 1, 2, 3]

        return [0, 1]

    def get_reward(self, action, done):
        player_hand_val = self.get_hand_value(self.player_hand)
        dealer_hand_val = self.get_hand_value(self.dealer_hand)
        if not done:
            return 0
        if player_hand_val == 21 and self.first_round:  # Natural blackjack
            return 1.5
        if action == 3:  # Surrender
            return 0.5
        if 21 >= player_hand_val > dealer_hand_val:
            return 1 * self.bet_multiplier
        if player_hand_val <= 21 and dealer_hand_val > 21:
            return 1 * self.bet_multiplier
        if player_hand_val > 21:
            return -1 * self.bet_multiplier
        if player_hand_val == dealer_hand_val:
            return 0
        if dealer_hand_val > 21:
            return 1 * self.bet_multiplier
        return -1 * self.bet_multiplier

    def deal_card(self) -> Card:
        card = Card(
            random.choice([face for _, (face, num_in_decks) in enumerate(self.deck.items()) if num_in_decks > 0]))
        self.deck[card.face] -= 1
        return card

    @staticmethod
    def get_hand_value(hand):
        value = 0
        aces = 0
        for card in hand:
            if card == "A":
                aces += 1
            else:
                value += card.value

        while aces:
            if aces == 1 and value <= 10:
                value += 11
            else:
                value += 1
            aces -= 1

        return value

    def dealer_plays(self):
        if self.get_hand_value(self.player_hand) > 21:
            return
        if self.get_hand_value(self.dealer_hand) == 21:
            return
        while self.get_hand_value(self.dealer_hand) <= 16:
            self.dealer_hand.append(self.deal_card())

    def is_busted(self):
        if self.get_hand_value(self.player_hand) > 21:
            return True

    def render(self):
        if self.player_round:
            print("Dealer hand: [" + self.dealer_hand[0].face + ", ?], Value: ?")
        else:
            print("Dealer hand: " + str(self.dealer_hand) + ", Value: " + str(self.get_hand_value(self.dealer_hand)))
        print("Player hand: " + str(self.player_hand) + ", Value: " + str(self.get_hand_value(self.player_hand)))


import torch
from torch.utils.data import Dataset
import torchvision
import torch.nn as nn

import random
from pathlib import Path
import os,sys
import math
import numpy as np
import json
import datasets.thumos_utils as data_utils

class ThumosDetection(Dataset):
    """
       thumos dataset
    """
    def __init__(self, vid_data_file, num_classes, actions_dict, features_root, sample_rate,split,out_size):
        self.vid_data_file = vid_data_file
        self.num_classes = num_classes
        self.actions_dict = actions_dict
        self.features_root = features_root
        self.sample_rate = sample_rate
        self.split = split
        self.out_size = out_size

        self.ids = self.read_data(self.vid_data_file,self.split)
        self.prepare = ThumosVideoData(self.sample_rate,self.actions_dict, self.split, self.out_size)

    def __len__(self):
        return len(self.ids)

    def __getitem__(self, idx):
        vid = self.ids[idx]
        target = {}
        target = {'video_id': vid, 'annotations': target}

        target['annotations']['feature_path'] = self.features_root + "/" + vid + '.npy'
        with open(self.vid_data_file,'r') as f:
            data = json.load(f)
        target['annotations']['actions'] = data[vid]['actions']
        target['annotations']['duration'] = data[vid]['duration']
        inp, target = self.prepare(target)
        return inp, target

    def read_data(self, vid_data_file,split):
        with open(vid_data_file,'r') as f:
            data = json.load(f)

        list_of_examples = []
        for k,v in data.items():
            if v['subset'] != split: continue
            if not os.path.exists(os.path.join(self.features_root,k+'.npy')): continue
            if os.path.exists(os.path.join(self.features_root,k+'.npy')):
                k_shape = np.load(os.path.join(self.features_root,k+'.npy')).shape[0]

            list_of_examples.append((k,k_shape))

        list_of_examples = sorted(list_of_examples, key=lambda x:x[1])
        list_of_examples = [e[0] for e in list_of_examples]

        return list_of_examples

class ThumosVideoData(object):
    def __init__(self,sample_rate,actions_dict,mode,out_size):
        self.sample_rate = sample_rate
        self.actions_dict = actions_dict
        self.mode = mode
        self.out_size = out_size
        pass

    def __call__(self, target):
        vid =  target['video_id']
        video_id = data_utils.getVideoId(vid)

        anno = target['annotations']

        feature = np.load(anno['feature_path'])

        triplets = anno['actions']
        duration = anno['duration']
        set_targets = {}
        set_targets['video_id'] = torch.tensor(video_id)
        set_targets['segments'] = []
        set_targets['labels'] = []

        for triplet in triplets:
            set_targets['segments'].append([triplet[1]/duration,triplet[2]/duration])
            set_targets['labels'].append(triplet[0])

        feature = torch.tensor(feature,dtype=torch.float32).squeeze_(1).squeeze_(1).permute(1,0)

        out_size = self.out_size
        offset = 32

        if feature.size(1) <= out_size:
            upsampled_size = feature.size(1) * 4
            initial_size = feature.size(1)
            tile_idx = np.zeros(initial_size)
            for i in range(initial_size):
                for j in range(math.floor(upsampled_size/initial_size)):
                    tile_idx[i] += 1
            for i in range(initial_size):
                tile_idx[i] += 1
                if np.sum(tile_idx) == upsampled_size: break
            feature = feature.repeat_interleave(torch.LongTensor(tile_idx),dim=1)

            if self.mode == 'training':
                fidx = sorted(random.sample(list(range(0,feature.size(1))),initial_size))
                feature = feature[:,torch.tensor(fidx)].squeeze(0)

            elif self.mode in ['testing', 'validation']:
                fidx = list(np.linspace(0,feature.size(1)-1,initial_size))
                fidx = [int(idx) for idx in fidx]
                feature = feature[:,torch.tensor(fidx)].squeeze(0)

        elif feature.size(1) > out_size:
            if self.mode == 'training':
                if len(list(range(0+offset,feature.size(1)-offset))) < out_size:
                    initial_size = len(list(range(0+offset,feature.size(1)-offset)))
                    tile_idx = np.zeros(initial_size)
                    for i in range(initial_size):
                        for j in range(math.floor(out_size/feature.size(1))):
                            tile_idx[i] += 1

                    for i in range(initial_size):
                        tile_idx[i] += 1
                        if np.sum(tile_idx) == out_size: break

                    sample_idx = torch.tensor(list(range(0+offset,feature.size(1)-offset))).repeat_interleave(torch.LongTensor(tile_idx))
                    fidx = sorted(sample_idx.tolist())
                else:
                    fidx = sorted(random.sample(list(range(0+offset,feature.size(1)-offset)),out_size))

                feature_sampled = []
                for idx in fidx:
                    feature_sampled.append(torch.mean(feature[:,torch.tensor(list(range(idx-offset,idx+offset)))],dim=1).unsqueeze(1))
                feature = torch.cat(feature_sampled,dim=1)

            elif self.mode in ['testing', 'validation']:
                fidx = list(np.linspace(0+offset,feature.size(1)-offset,out_size))
                fidx = [int(idx) for idx in fidx]
                feature_sampled = []
                for idx in fidx:
                    feature_sampled.append(torch.mean(feature[:,torch.tensor(list(range(idx-offset,idx+offset)))],dim=1).unsqueeze(1))
                feature = torch.cat(feature_sampled,dim=1)

        set_targets['num_feat'] = torch.tensor(feature.size(1),dtype = torch.float32)
        set_targets['segments'] = torch.tensor(set_targets['segments'],dtype=torch.float32)
        set_targets['labels'] = torch.tensor(set_targets['labels'],dtype=torch.int64)
        set_targets['length'] = torch.tensor([duration/self.sample_rate],dtype=torch.float32)
        set_targets['sample_rate'] = torch.tensor([self.sample_rate],dtype=torch.float32)
        set_targets['fps'] = torch.tensor([float(30)/self.sample_rate],dtype=torch.float32)
        set_targets['num_classes'] = torch.tensor([len(list(self.actions_dict.keys()))],dtype=torch.float32)
        return feature, set_targets

def build_thumos_detection(mode, args):
    root = Path(args.data_root)
    assert root.exists()

    features_path = os.path.join(args.data_root, args.features)
    vid_data_file = os.path.join(args.data_root, args.dataset+ "_action.json")
    # action category mappings
    mapping_file = os.path.join(args.data_root, "action_mapping.txt")
    file_ptr = open(mapping_file, 'r')
    actions = file_ptr.read().split('\n')[:-1]
    file_ptr.close()
    actions_dict = dict()
    for a in actions:
        actions_dict[a.split()[1]] = int(a.split()[0])

    num_classes = len(actions_dict)
    assert num_classes == args.num_classes   # check if the number of classes are determined correctly

    dataset = ThumosDetection(vid_data_file, args.num_classes, actions_dict, features_path, args.sample_rate, mode, args.num_inputs)

    return dataset