from typing import Optional

import torch
import torch.nn as nn
from torch import Tensor

import detectron2.layers as L
from detectron2.config import configurable
from detectron2.structures import Boxes, Instances
from detectron2.utils.registry import Registry
from pytorch3d.ops import corresponding_points_alignment
from pytorch3d.transforms import euler_angles_to_matrix

from nocpred.structures import (
    DepthPoints,
    MeshGrids,
    NOCs,
    RotationMats,
    Rotations,
    Scales,
    transform_grid,
    Translations,
)

ROI_NOC_HEAD_REGISTRY = Registry('ROI_NOC_HEAD')

@ROI_NOC_HEAD_REGISTRY.register()
class ConvUpsampleNOCHead(nn.Module):
    @configurable
    def __init__(
        self,
        input_size: int = 256,
        hidden_size: int = 256,
        num_pre_convs: int = 4,
        num_post_convs: int = 1,
        num_mlp_hiddens: int = 2,
        filter_size: int = 3,
        num_cond: int = 0,
        num_mlp_cond: int = 0,
        use_deconv: bool = False,
        use_pixel_shuffle: bool = True,
        output_size: int = 3,
        noc_loss_weight: float = 1.,
        proc_loss: bool = False,
        proc_rot_weight: float = 1.,
        proc_trans_weight: float = 1.,
        min_proc: int = 10,
        eval_mask_thresh: float = 0.5,
        predict_scale: bool = False,
        predict_alignment: bool = False,
        num_classes: int = 9,
        pooler_size: tuple[int, int] = (16, 16),
    ):
        super().__init__()
        padding = filter_size // 2 - (1 - filter_size % 2)

        self.pre_convs = nn.Sequential(*[
            L.Conv2d(
                input_size if i == 0 else hidden_size,
                hidden_size,
                filter_size,
                padding=padding,
                activation=nn.ReLU(True),
            )
            for i in range(num_pre_convs)
        ])

        if use_deconv:
            self.upsample = nn.Sequential(
                L.ConvTranspose2d(hidden_size, hidden_size, 2, stride=2),
                nn.ReLU(True),
            )
        elif use_pixel_shuffle:
            self.upsample = nn.Sequential(
                L.Conv2d(
                    hidden_size,
                    hidden_size * 4,
                    # 1,
                    filter_size,
                    padding=padding,
                    activation=nn.ReLU(True),
                ),
                nn.PixelShuffle(2),
            )
        else:
            self.upsample = nn.Upsample(scale_factor=2, mode='nearest')

        self.post_convs = nn.Sequential(*[
            L.Conv2d(
                hidden_size + (i == 0) * num_cond,
                hidden_size,
                filter_size,
                padding=padding,
                activation=nn.ReLU(True),
            )
            for i in range(num_post_convs)
        ])

        self.shared_mlp = nn.Sequential(*[
            L.Conv2d(
                hidden_size + (i == 0) * num_mlp_cond,
                hidden_size,
                1,
                activation=nn.ReLU(True),
            )
            for i in range(num_mlp_hiddens)
        ])

        self.output = L.Conv2d(hidden_size, output_size, 1)

        self.noc_loss_weight = noc_loss_weight
        self.loss = nn.L1Loss(reduction='none')

        self.proc_loss = proc_loss and (proc_rot_weight > 0 or proc_trans_weight > 0)
        self.proc_rot_weight = proc_rot_weight
        self.proc_trans_weight = proc_trans_weight
        self.min_proc = min_proc
        self.eval_mask_thresh = eval_mask_thresh

        if self.proc_loss:
            self.rot_loss = nn.L1Loss()
            self.trans_loss = nn.HuberLoss()

        self.predict_alignmnet = predict_alignment
        self.predict_scale = predict_scale
        mlp_input_size = (pooler_size[0] // 2) * (pooler_size[1] // 2) * hidden_size
        self.num_classes = num_classes
        if predict_scale:
            self.scale_head = self._make_mlp(
                input_size, mlp_input_size, num_classes, hidden_size, 3
            )
            self.scale_loss = nn.L1Loss()

        self.sym_head = self._make_mlp(input_size, mlp_input_size, num_classes, hidden_size, 4)
        self.sym_loss = nn.CrossEntropyLoss(torch.as_tensor([.22, .43, .78, 1.]))

    @classmethod
    def from_config(cls, cfg, shape_spec: L.ShapeSpec):
        return {
            'input_size': shape_spec.channels,
            'noc_loss_weight': cfg.MODEL.ROI_NOC_HEAD.NOC_LOSS_WEIGHT,
            'proc_loss': cfg.MODEL.ROI_NOC_HEAD.PROCRUSTES_LOSS,
            'proc_rot_weight': cfg.MODEL.ROI_NOC_HEAD.PROCRUSTES_ROTATION_WEIGHT,
            'proc_trans_weight': cfg.MODEL.ROI_NOC_HEAD.PROCRUSTES_TRANSLATION_WEIGHT,
            'eval_mask_thresh': cfg.TEST.NOC_MASK_THRESH,
            'use_deconv': cfg.MODEL.ROI_NOC_HEAD.USE_DECONV,
            'use_pixel_shuffle': cfg.MODEL.ROI_NOC_HEAD.USE_PIXEL_SHUFFLE,
            'predict_scale': cfg.MODEL.ROI_NOC_HEAD.PREDICT_SCALE,
            'pooler_size': (shape_spec.width, shape_spec.height),
            'num_classes': cfg.MODEL.ROI_HEADS.NUM_CLASSES,
        }

    @staticmethod
    def _make_mlp(
        input_size: int,
        mlp_input_size: int,
        num_classes: int,
        hidden_size: int,
        output_size: int,
        mlp_hidden_size: int = 1024,
    ):
        return nn.Sequential(
            L.Conv2d(
                input_size,
                hidden_size,
                (5, 5),
                stride=2,
                padding=2,
                activation=nn.ReLU(True),
            ),
            nn.Flatten(),
            L.Linear(mlp_input_size, mlp_hidden_size),
            nn.ReLU(True),
            L.Linear(mlp_hidden_size, mlp_hidden_size),
            nn.ReLU(True),
            L.Linear(mlp_hidden_size, output_size * num_classes),
        )

    @property
    def device(self):
        return next(self.parameters()).device

    def forward(self, features: Tensor, instances: list[Instances]):
        nocs, scales, syms = self.run_network(features)
        if self.training:
            return self.compute_losses(instances, nocs, scales, syms)
        else:
            return self.inference(instances, features, nocs, scales, syms)

    def inference(
        self,
        instances: list[Instances],
        features: Tensor,
        nocs: Tensor,
        scales: Optional[Tensor],
        syms: Tensor,
    ) -> list[Instances]:
        if features.numel() == 0:
            # Safe-guard against no proposals case
            for instance in instances:
                instance.pred_nocs = NOCs(tensor=torch.zeros(
                    (0, 3, features.size(-2) * 2, features.size(-1) * 2),
                    device=self.device,
                ))
                instance.pred_xy_grids = MeshGrids(tensor=torch.zeros(
                    (0, 2, features.size(-2) * 2, features.size(-1) * 2),
                    device=self.device,
                ))
                if self.predict_scale:
                    instance.pred_scales = Scales(tensor=torch.ones(0, 3, device=self.device))

                instance.pred_syms = torch.zeros(0, 4, dtype=torch.long, device=self.device)

        else:
            # NOTE: This can also be enabled for training if NOCs are needed outside
            lens = list(map(len, instances))
            for noc, instance in zip(nocs.split(lens), instances):
                grid = MeshGrids(instance.image_size[::-1], nocs.size(0), device=self.device)\
                    .crop_and_resize(
                        instance.pred_boxes.tensor,
                        crop_size=noc.size(-1),
                        use_interpolate=True,
                        wrap_output=False,
                )
                raw_mask = (instance.pred_masks > self.eval_mask_thresh)
                instance.pred_xy_grids = MeshGrids(tensor=(grid * raw_mask))
                instance.pred_nocs = NOCs(tensor=(noc * raw_mask))

            if self.predict_scale:
                for scale, instance in zip(scales.split(lens), instances):
                    classes = instance.pred_classes
                    scale = scale[torch.arange(classes.numel()), classes]
                    instance.pred_scales = Scales(tensor=scale)

            for sym, instance in zip(syms.split(lens), instances):
                classes = instance.pred_classes
                sym = sym[torch.arange(classes.numel()), classes]
                instance.pred_syms = sym.argmax(-1)

        return instances

    def run_network(self, x: Tensor) -> tuple[Tensor, Optional[Tensor]]:
        nocs = self.output(self.shared_mlp(self.post_convs(self.upsample(self.pre_convs(x)))))
        scales = None
        if self.predict_scale:
            scales = self.scale_head(x).reshape(-1, self.num_classes, 3)
        syms = self.sym_head(x).reshape(-1, self.num_classes, 4)
        return nocs, scales, syms

    def compute_losses(
        self,
        instances: list[Instances],
        nocs: Tensor,
        scales: Optional[Tensor],
        syms: Tensor,
    ) -> dict[str, Tensor]:

        assert self.training

        losses = {}

        boxes = Boxes.cat([p.proposal_boxes for p in instances])
        gt_nocs = NOCs.cat([p.gt_nocs for p in instances])
        valid_mask = L.cat([p.gt_valid_nocs for p in instances])
        if not valid_mask.all():
            nocs = nocs[valid_mask]
            boxes = boxes[valid_mask]
            gt_nocs = gt_nocs[valid_mask]

        gt_classes = L.cat([p.gt_classes for p in instances])
        syms = syms[torch.arange(gt_classes.numel()), gt_classes]
        gt_syms = L.cat([p.gt_syms for p in instances])
        losses['loss_sym'] = self.sym_loss(syms, gt_syms)

        if len(boxes) == 0:  # No valid NOCs!
            loss_noc = torch.tensor(0., device=gt_nocs.device)
        else:
            gt_nocs = gt_nocs.crop_and_resize_with_grids_from_boxes(boxes, nocs.size(-1))
            # loss_masks = gt_nocs.masks()
            # diffs: Tensor = self.loss(nocs, gt_nocs.tensor) * loss_masks
            loss_noc = self._sym_aware_noc_loss(gt_syms[valid_mask], nocs, gt_nocs)
            # loss_noc = diffs.sum() / loss_masks.sum().float().clamp(1e-5)
            # self._debug(instances, loss_masks, valid_mask, boxes, gt_nocs)
        losses['loss_noc'] = self.noc_loss_weight * loss_noc

        if self.predict_scale:
            scales = scales[torch.arange(gt_classes.numel()), gt_classes]
            gt_scales = Scales.cat([p.gt_scales for p in instances]).tensor
            losses['loss_scale'] = self.scale_loss(scales, gt_scales)

        if self.proc_loss:
            losses.update(self._compute_proc_losses(instances, nocs, boxes, valid_mask, scales))

        return losses

    def _compute_proc_losses(
        self,
        instances: list[Instances],
        nocs: Tensor,
        boxes: Boxes,
        valid_mask: Tensor,
        scales: Optional[Tensor] = None,
    ) -> dict[str, Tensor]:

        device = nocs.device
        empty_output = {
            'loss_proc_rot': torch.tensor(0., device=device),
            'loss_proc_trans': torch.tensor(0., device=device),
        }

        # Handle emptry case
        if len(boxes) == 0:
            return empty_output

        # Get gt depths
        depths = DepthPoints.cat([p.gt_depth_points for p in instances])[valid_mask]
        depths = depths.crop_and_resize_with_grids_from_boxes(boxes, nocs.size(-1))
        masks = depths.masks()

        # Remove invalids
        new_valid_mask = (masks.sum((1, 2, 3)) >= self.min_proc)
        if not new_valid_mask.any():
            return empty_output

        # Prepare Procrustes inputs
        # if scales is None:
        scales = Scales.cat([p.gt_scales for p in instances])
        # else:
        #    scales = Scales(scales)
        scales = scales[valid_mask][new_valid_mask]
        nocs = transform_grid(nocs[new_valid_mask], scales=scales)

        masks = masks[new_valid_mask]

        poses = L.cat([p.gt_poses for p in instances])[valid_mask][new_valid_mask]
        depths = depths.tensor[new_valid_mask]
        depths = transform_grid(
            depths,
            masks=masks,
            translations=Translations(poses[:, :3, 3]),
            rotations=RotationMats(poses[:, :3, :3]),
        )

        dtype = nocs.dtype
        nocs = nocs.flatten(2).transpose(1, 2).double()
        depths = depths.flatten(2).transpose(1, 2).double()
        masks = masks.flatten(1).double()

        # Perform the Procrustes alignment
        so3 = corresponding_points_alignment(nocs, depths, masks)
        rots, transes = so3.R.transpose(1, 2).to(dtype), so3.T.to(dtype)

        # Compute procrustes loss
        gt_rots = Rotations.cat([p.gt_rotations for p in instances])
        gt_rots = gt_rots[valid_mask][new_valid_mask].as_rotation_matrices().tensor
        gt_transes = Translations.cat([p.gt_translations for p in instances])
        gt_transes = gt_transes.tensor[valid_mask][new_valid_mask]

        # loss_rot = self.rot_loss(rots.flatten(1), gt_rots)
        syms = L.cat([p.gt_syms for p in instances])[valid_mask][new_valid_mask]
        loss_rot = self._sym_aware_rot_loss(syms, rots, gt_rots)
        loss_trans = self.trans_loss(transes, gt_transes)
        return {
            'loss_proc_rot': self.proc_rot_weight * loss_rot,
            'loss_proc_trans': self.proc_trans_weight * loss_trans,
        }

    def _sym_aware_rot_loss(self, gt_syms: Tensor, rots: Tensor, gt_rots: Tensor):
        rots, gt_rots = rots.flatten(1), gt_rots.flatten(1)

        '''if torch.all(gt_syms == 0):
            return self.rot_loss(rots, gt_rots)'''

        losses = []
        rots, gt_rots = rots.cpu().double(), gt_rots.cpu().double()
        gt_syms = gt_syms.cpu().tolist()
        for i, sym_type in enumerate(gt_syms):
            loss = self.rot_loss(rots[i], gt_rots[i])
            angles = self._get_angles_for_sym(sym_type)
            for angle in angles:
                rots_ = rots[i].reshape(3, 3) @ euler_angles_to_matrix(
                    torch.as_tensor([0., angle, 0.])[None].double(), 'XYZ'
                ).reshape(3, 3)
                loss = torch.minimum(loss, self.rot_loss(rots_.reshape_as(gt_rots[i]), gt_rots[i]))
            losses.append(loss)
        return torch.mean(torch.stack(losses)).float().to(self.device)

    def _sym_aware_noc_loss(self, syms: Tensor, nocs: Tensor, gt_nocs_wrap: NOCs):
        masks = gt_nocs_wrap.masks()
        gt_nocs = gt_nocs_wrap.tensor
        base_losses = self._compute_noc_losses(nocs, gt_nocs, masks)
        losses = [base_losses[syms == 0]]
        for i in range(1, 4):
            sym_mask = syms == i
            if not sym_mask.any():
                continue
            min_losses = base_losses[sym_mask]
            angles = self._get_angles_for_sym(i, 6)
            nocs_, gt_nocs_ = nocs[sym_mask], gt_nocs[sym_mask]
            masks_ = masks[sym_mask]
            for angle in angles:
                mat = euler_angles_to_matrix(
                    torch.as_tensor([0., angle, 0.], device=self.device)[None], 'XYZ'
                )
                nocs_a = transform_grid(nocs_, rotations=RotationMats(mat))
                min_losses = torch.minimum(
                    min_losses,
                    self._compute_noc_losses(nocs_a, gt_nocs_, masks_),
                )
            losses.append(min_losses)
        return torch.cat(losses).sum() / masks.sum().clamp(1e-5)

    def _compute_noc_losses(self, nocs: Tensor, gt_nocs: Tensor, masks: Tensor):
        base_diffs: Tensor = self.loss(nocs, gt_nocs) * masks
        return base_diffs.sum((1, 2, 3))

    @staticmethod
    def _get_angles_for_sym(sym_type: int, bins_for_inf: int = 36) -> list[float]:
        if sym_type == 1:
            angles = [torch.pi]
        elif sym_type == 2:
            angles = [torch.pi / 2, torch.pi, 3 * torch.pi / 2]
        elif sym_type == 3:
            angles = [x * (2 * torch.pi / bins_for_inf) for x in range(1, bins_for_inf)]
        else:
            angles = []
        return angles

    def _debug(self, instances, loss_masks, valid_mask, boxes, gt_nocs):
        gt_depths = DepthPoints.cat([p.gt_depth_points for p in instances])[valid_mask]
        gt_depths = gt_depths.crop_and_resize_with_grids_from_boxes(boxes, gt_nocs.image_size[-1])
        pose = L.cat([p.gt_poses for p in instances])[valid_mask]
        gt_depths = pose[..., :3, :3] @ gt_depths.tensor.flatten(2) + pose[..., :3, 3:4]
        gt_depths *= loss_masks.flatten(2)

        gt_translations = Translations.cat([p.gt_translations for p in instances])
        gt_rotations = Rotations.cat([p.gt_rotations for p in instances])
        gt_scales = Scales.cat([p.gt_scales for p in instances])
        gt_nocs = transform_grid(
            gt_nocs.tensor, loss_masks, gt_translations, gt_rotations, gt_scales
        )
        print((gt_nocs - gt_depths.view_as(gt_nocs))[loss_masks.expand_as(gt_nocs)].mean())
        from IPython import embed; embed()  # noqa

def build_noc_head(cfg, input_shape: L.ShapeSpec) -> nn.Module:
    return ROI_NOC_HEAD_REGISTRY.get(cfg.MODEL.ROI_NOC_HEAD.NAME)(cfg, input_shape)


import sys
import socket
import select
from logzero import logger

from pprint import pprint

class Forwarder:
    def __init__(self, src_ip, src_port, dst_ip, dst_port) -> None:
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

        logger.info(f"Binding to {src_ip}:{src_port}")
        self.sock.bind((src_ip, src_port))
        self.sock.listen(5)
        self.target = (dst_ip, dst_port)

    def coding(self, data):
        return data

    def exchange_loop(self, client, remote):
        while True:
            r, w, e = select.select([client, remote], [], [])

            if client in r:
                data = client.recv(4096)
                logger.debug(f" CLIENT > REMOTE : {len(data)} bytes")
                pprint(self.coding(data))
                if remote.send(self.coding(data)) <= 0:
                    return

            if remote in r:
                data = remote.recv(4096)
                logger.debug(f" CLIENT < REMOTE : {len(data)} bytes")
                pprint(self.coding(data))
                if client.send(self.coding(data)) <= 0:
                    return

    def run(self):
        remote = None
        try:
            while True:
                client, addr = self.sock.accept()
                logger.info(f"[NEW] CLIENT({addr}) forward REMOTE({self.target}) ")

                remote = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                remote.connect(self.target)

                self.exchange_loop(client, remote)

                client.close()
                logger.info(f"[CLOSE] CLIENT({addr[0]})")

        except Exception as e:
            logger.error(e)
        finally:
            if remote is not None:
                remote.close()
            logger.info(f"[CLOSE] REMOTE({self.target[0]})")

class SecureForwarderXD(Forwarder):

    def xor(self, data, key_byte=42):
        result = b""
        for b in data:
            result += bytes([b ^ key_byte])
        return result

    def coding(self, data):
        return self.xor(data)

def main() -> None:
    print(f"Hello main from : {__file__}")

    src_ip, src_port = sys.argv[1].split(":")
    src_port = int(src_port)
    dst_ip, dst_port = sys.argv[2].split(":")
    dst_port = int(dst_port)

    logger.info(f"TCP Forward {src_ip}:{src_port} ===> {dst_ip}:{dst_port}")
    if len(sys.argv) > 3 and sys.argv[3] == "secure":
        logger.info("Creation of \"SECURE\" forwarder")
        proxy = SecureForwarderXD(src_ip, src_port, dst_ip, dst_port)
    else:
        logger.info("Creation of NORMAL forwarder")
        proxy = Forwarder(src_ip, src_port, dst_ip, dst_port)
    proxy.run()

if __name__ == "__main__":
    main()


from __future__ import division
from __future__ import absolute_import
from __future__ import unicode_literals
from builtins import range
import argparse
import os
os.environ['OPL_NUM_THREADS'] = "1"
os.environ['MKL_NUM_THREADS'] = "1"
os.environ['NUMEXPR_NUM_THREADS'] = "1"
import numpy as np
import pandas as pd
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
import os.path as osp
import sys

from .simulate import get_parameters
from . import util as ut
from . import newick
from collections import OrderedDict, defaultdict

def get_hatch(paramtype):
    hatch = None
    if paramtype == 'bottleneck':
        hatch = 'o'
    elif paramtype == 'rate':
        hatch = '\\'
    return hatch

def get_parameter_names(filename):
    with open(filename) as fin:
        tree_str = fin.read().strip()
    branch_names = []
    leaf_names = []
    tree = newick.loads(tree_str, look_for_multiplier = True,
                       length_parser = ut.length_parser_str)[0]
    vset = set([])
    for node in tree.walk('postorder'):
        if node == tree:
            continue
        vset.add(node.varname)
    vnames = sorted(list(vset))
    param_names = ['loglike']
    for vname in vnames:
        param_names.append(vname + '_l')
    for vname in vnames:
        param_names.append(vname + '_m')
    param_names.append('root')
    param_names.append('ppoly')
    return param_names

def get_chains(dat, num_walkers):
    chains = []
    for i in range(num_walkers):
        chains.append(dat.iloc[i::num_walkers,:].copy())
    return chains

def get_acceptance(chain):
    assert 3/2 == 1.5
    rows0 = chain.iloc[:-1,:].values
    rows1 = chain.iloc[1:,:].values
    num_changes = np.sum(np.all(~np.equal(rows0, rows1), axis = 1))
    num_comparisons = rows0.shape[0]
    return num_changes / num_comparisons

def get_len_limits(trees, ages_file, lower_drift, upper_drift,
        lower_bottleneck = 2, upper_bottleneck = 500):
    is_bottleneck = {}
    ages_dat = pd.read_csv(ages_file, sep = '\t')

    # first, get min and max multipliers
    mins = defaultdict(lambda: np.inf)
    maxes = defaultdict(lambda: -np.inf)
    for tree in trees:
        for node in tree.walk():
            varname = node.varname
            is_bottleneck[varname] = node.is_bottleneck
            mult = node.multipliername
            if mult is None:
                mins[varname] = 1.0
                maxes[varname] = 1.0
            else:
                minv, maxv = ages_dat[mult].min(), ages_dat[mult].max()
                if varname in mins:
                    mins[varname] = min(mins[varname], minv)
                else:
                    mins[varname] = minv
                if varname in maxes:
                    maxes[varname] = max(maxes[varname], maxv)
                else:
                    maxes[varname] = maxv
    limits = {}

    # so far, it has just been the multipliers, now need to make them limits
    assert set(mins.keys()) == set(maxes.keys())
    for varname in list(mins.keys()):
        minv = mins[varname]
        maxv = maxes[varname]
        ld = np.array([minv, maxv])
        #mults = limits[varname]
        if is_bottleneck[varname]:
            ld, ud = lower_bottleneck, upper_bottleneck
        else:
            ld, ud = lower_drift, upper_drift
        limits[varname] = [np.log10(ld / minv),
                np.log10(ud / maxv)]
    return limits

def make_figures(
        results, tree_files, num_walkers, ages_file, prefix = None,
        img_format = 'png', dpi = 300, colors_and_types = None,
        length_prior_limits = (-6,0.477), mutation_prior_limits = (-8,-1),
        true_parameters = None, which_plot = 'histograms',
        trace_burnin_steps = 500, posterior_burnin_frac = 0.3,
        add_title = True, log_uniform_drift_priors = False):

    do_histograms = False
    if which_plot in ('histograms', 'both'):
        do_histograms = True

    do_traces = False
    if which_plot in ('traces', 'both'):
        do_traces = True

    do_corner = False
    if which_plot == 'corner':
        do_corner = True

    validparamtypes = ('fixed', 'rate', 'bottleneck')

    ###################
    # output filenames
    ###################

    if prefix is None:
        path, filename = osp.split(results)
        prefix = filename.split('.')[0]

    trace_file = 'traces_' + prefix + '.' + img_format
    posterior_histograms_file = ('posterior_histograms_' + prefix + '.' +
            img_format)
    corner_file = ('corner_' + prefix + '.' + img_format)

    # loading colors and (drift) parameter types
    if do_histograms:
        if colors_and_types is not None:
            colors = OrderedDict()
            types = OrderedDict()
            #priortypes = OrderedDict()
            with open(colors_and_types) as inp:
                for line in inp:
                    spline = line.split()
                    param, color, paramtype = spline[:3]
                    if paramtype not in validparamtypes:
                        raise ValueError('invalid parameter type: must be "fixed", '
                                         '"rate", or "bottleneck"')
                    colors[param] = color
                    types[param] = paramtype

    ###################
    # loading the data
    ###################
    all_varnames = []
    trees = []
    with open(tree_files) as trees_in:
        for line in trees_in:
            fn = line.strip()
            with open(fn) as fin:
                tree_str = fin.read().strip()
                tree = newick.loads(tree_str,
                                    look_for_multiplier = True,
                                    length_parser = ut.length_parser_str)[0]
            trees.append(tree)
            cols = get_parameter_names(fn)
            varnames = [col[:-2] for col in cols if col.endswith('_l')]
            all_varnames += varnames
    try:
        # first try no header
        dat_m1 = pd.read_csv(results,
                sep = '\t', header = None, names = cols, dtype = np.float64,
                comment = '#')
    except ValueError:
        # then with header
        dat_m1 = pd.read_csv(results,
                sep = '\t', header = 0, names = cols, dtype = np.float64,
                comment = '#')
    # drift parameters are output in natural scale
    #dat_m1.loc[:,dat_m1.columns.str.contains('_l')] = np.log10(np.abs(
    #        dat_m1.loc[:,dat_m1.columns.str.contains('_l')]))
    # mutation parameters are given in log-10 scale
    dat_m1.loc[:,dat_m1.columns.str.contains('_m')] = -np.abs(
            dat_m1.loc[:,dat_m1.columns.str.contains('_m')])
    dat_m1.loc[:,'root'] = -np.abs(dat_m1.loc[:,'root'])
    dat_m1.loc[:,'ppoly'] = -np.abs(dat_m1.loc[:,'ppoly'])
    num_walkers = num_walkers
    burnin_steps = trace_burnin_steps
    chains = get_chains(dat_m1.iloc[burnin_steps*num_walkers:,:].copy(),
            num_walkers)

    # getting (optional) length prior limits
    if do_histograms:
        if length_prior_limits is not None:
            lower_l, upper_l = length_prior_limits
            lower_l = 10**float(lower_l)
            upper_l = 10**float(upper_l)
            length_limits = get_len_limits(trees, ages_file, lower_l, upper_l)

    # getting true parameters
    if true_parameters is not None:
        true_bls, true_mrs, true_ab, true_ppoly = get_parameters(
                true_parameters, tree)
        true_params = [true_bls[v] for v in varnames]
        true_params += [true_mrs[v] for v in varnames]
        true_params += [true_ab, true_ppoly]

    if do_traces:
        f, axes = plt.subplots(
                dat_m1.shape[1], 1, figsize = (10, 2.5*dat_m1.shape[1]))
        for i in range(dat_m1.shape[1]):
            col = list(dat_m1.columns)[i]
            ax = axes[i]
            if not np.any(dat_m1[col] <= 0.0):
                for chain in chains:
                    ax.plot(np.arange(chain.shape[0]),np.log10(chain.iloc[:,i]),
                            lw = 0.02)
                plot_log = True
            else:
                for chain in chains:
                    ax.plot(np.arange(chain.shape[0]),(chain.iloc[:,i]), lw = 0.02)
                plot_log = False
            if true_parameters is not None and i > 0:
                true_p = true_params[i-1]
                if plot_log:
                    ax.axhline(np.log10(true_p))
                else:
                    # now all the parameters are log10 transformed
                    ax.axhline(np.log10(true_p))
            ax.set_title(dat_m1.columns[i])
        f.tight_layout()
        plt.savefig(trace_file)

    # posterior histograms
    if do_histograms:
        if colors_and_types is None:
            final = False  # not a 'final', publication plot
            num_plots = dat_m1.shape[1]
            nrows = 4
            ncols = int(np.ceil(num_plots / nrows))
            f, axes = plt.subplots(nrows,ncols, figsize = (12,12*nrows/ncols))

        else:
            final = True   # a 'final', publication plot
            nrows = 2
            ncols = len(varnames)
            num_plots = (dat_m1.shape[1]-3)*2  # (no loglike, or the two root params)
            f, axes = plt.subplots(nrows,ncols, figsize = (2*ncols,4))
            mpl.rcParams.update({'font.size': 5})
            mpl.rcParams.update({'axes.labelsize': 'small'})

        burnin = int(np.floor(dat_m1.shape[0] * posterior_burnin_frac))
        dat_burn = dat_m1.iloc[burnin:,:]
        plotted = np.tile([False], nrows*ncols)
        plotted.shape = (nrows,ncols)

        if final:
            if not add_title:
                mpl.rc('text', usetex = False)
            #mpl.rcParams.update({'font.family': 'Arial'})
            colornames = list(colors.keys())
            for counter, var in enumerate(colornames):
                if true_parameters is not None:
                    true_l = np.log10(true_bls[var])
                    true_m = np.log10(true_mrs[var])
                # length
                if colors_and_types is not None:
                    try:
                        c = colors[var]
                        t = types[var]
                    except KeyError:
                        c = 'gray'
                        t = 'fixed'
                else:
                    c = 'gray'
                    t = 'fixed'
                hatch = get_hatch(t)

                ax = axes[0][counter]
                col = var + '_l'
                if length_prior_limits is not None:
                    minv, maxv = length_limits[var]
                else:
                    minv, maxv = dat_burn[col].min(), dat_burn[col].max()

                lower = np.floor(minv)
                upper = np.ceil(maxv)
                bins = np.linspace(lower, upper, 50)
                ticks = np.arange(lower, upper+1).astype(np.int)
                if ticks.shape[0] > 5:
                    ticks = np.arange(lower, upper+1, 2).astype(np.int)

                if length_prior_limits is not None:
                    x = np.linspace(minv, maxv, 1000)
                    if log_uniform_drift_priors:
                        y = 1.0/(maxv-minv)
                    else:
                        y = 10**x * np.log(10) / (10**maxv - 10**minv)
                    ax.fill_between(x, 0, y, alpha = 0.3, color = 'gray')

                ax.hist(dat_burn[col], bins, color = c, hatch = hatch,
                        normed = True)
                ax.set_xticks(ticks)
                #ax.set_xscale('log')
                ax.get_yaxis().set_visible(False)
                if true_parameters is not None:
                    ax.axvline(true_l, ls = '--', c = 'black', lw = 2)
                if add_title:
                    ax.set_title(col)

                ax = axes[1][counter]
                col = var + '_m'
                if mutation_prior_limits is None:
                    minv, maxv = dat_burn[col].min(), dat_burn[col].max()
                else:
                    minv, maxv = mutation_prior_limits
                    minv = float(minv)
                    maxv = float(maxv)
                lower = np.floor(minv)
                upper = np.ceil(maxv)
                ticks = np.arange(lower, upper+1).astype(np.int)
                if ticks.shape[0] > 5:
                    ticks = np.arange(lower, upper+1, 2).astype(np.int)
                bins = np.linspace(lower, upper, 50)

                if mutation_prior_limits is not None:
                    x = np.linspace(minv, maxv, 1000)
                    y = 1.0/(maxv-minv)
                    ax.fill_between(x, 0, y, alpha = 0.3, color = 'gray')

                ax.hist(dat_burn[col], bins, color = c, hatch = hatch,
                        normed = True)
                ax.set_xticks(ticks)
                ax.get_yaxis().set_visible(False)
                if true_parameters is not None:
                    ax.axvline(true_m, ls = '--', c = 'black', lw = 2)
                if add_title:
                    ax.set_title(col)

        else:  # not final
            counter = 0
            for row_idx in range(nrows):
                for col_idx in range(ncols):
                    ax = axes[row_idx][col_idx]
                    try:
                        col = cols[counter]
                    except IndexError:
                        break
                    if true_parameters is not None and counter > 0:
                        true_p = true_params[counter-1]
                    minv, maxv = dat_burn[col].min(), dat_burn[col].max()
                    if colors_and_types is not None:
                        var = col[:-2]
                        try:
                            c = colors[var]
                            t = types[var]
                        except KeyError:
                            c = 'gray'
                            t = 'fixed'
                    else:
                        c = 'gray'
                        t = 'fixed'
                    hatch = get_hatch(t)
                    if minv >= 0:
                        lower = np.floor(np.log10(minv))
                        upper = np.ceil(np.log10(maxv))
                        bins = np.logspace(lower, upper, 50)
                        ax.hist(dat_burn[col], bins, color = c, hatch = hatch)
                        ax.set_xscale('log')
                    else:
                        if col != 'loglike':
                            lower = np.floor(minv)
                            upper = np.ceil(maxv)
                            bins = np.logspace(lower, upper, 50)
                            datp = 10**dat_burn[col]
                            ax.hist(datp, bins, color = 'gray')
                            ax.set_xscale('log')
                        else:
                            lower = np.floor(minv)
                            upper = np.ceil(maxv)
                            bins = np.linspace(lower, upper, 50)
                            ax.hist(dat_burn[col], bins, color = 'gray')
                    if true_parameters is not None and counter > 0:
                        ax.axvline(true_p)
                    ax.set_title(col)
                    ax.get_yaxis().set_visible(False)
                    counter += 1
                    plotted[row_idx][col_idx] = True

            for row_idx in range(nrows):
                for col_idx in range(ncols):
                    if not plotted[row_idx][col_idx]:
                        ax = axes[row_idx][col_idx]
                        f.delaxes(ax)

        f.tight_layout()
        plt.savefig(posterior_histograms_file, dpi = dpi)

    if do_corner:
        try:
            import corner
        except ImportError:
            err = 'plotting corner plots requires installing corner'
            raise ImportError(err)

        try:
            import numpy.random as npr
        except ImportError:
            err = 'corner plots require numpy'
            raise ImportError(err)

        burnin = int(posterior_burnin_frac * dat_m1.shape[0] + 0.5)
        dat_burn = dat_m1.iloc[burnin:,:]
        sample_size = min(500000, dat_burn.shape[0])
        idxs = np.sort(npr.choice(dat_burn.shape[0], size = sample_size,
            replace = False))
        dat_samp = dat_burn.iloc[idxs,:]
        corner.corner(dat_samp)
        plt.savefig(corner_file)

def _run_make_figures(args):

    make_figures(
            results = args.results,
            tree_files = args.tree,
            num_walkers = args.num_walkers,
            ages_file = args.ages_file,
            prefix = args.prefix,
            img_format = args.format,
            dpi = args.dpi,
            colors_and_types = args.colors_and_types,
            length_prior_limits = args.length_prior_limits,
            mutation_prior_limits = args.mutation_prior_limits,
            true_parameters = args.true_parameters,
            which_plot = args.plot,
            trace_burnin_steps = args.trace_burnin_steps,
            posterior_burnin_frac = args.posterior_burnin_frac,
            add_title = args.add_title,
            log_uniform_drift_priors = not args.uniform_drift_priors)


import sqlite3
from aifc import Error

import flask
from flask import request, Flask

app = Flask(__name__)

@app.route('/')
def hello_world82():
    return 'Welcome to Smart Groceries'

@app.route('/shoppinglist')
def hello_world2():
    conn = create_connection("datenbank.db")
    return select_all_tasks(conn)

@app.route('/products')
def hello_world832():
    conn = create_connection("datenbank.db")
    chosen = request.args.get("Produktname")
    cur = conn.cursor()
    cur.execute("SELECT * FROM produktinformationen WHERE Produktname = ?", (chosen,))

    row = cur.fetchone()

    return flask.jsonify(row)

def select_all_tasks(conn):
    cur = conn.cursor()
    cur.execute("SELECT * FROM produktinformationen;")

    rows = cur.fetchall()

    return flask.jsonify(rows)

def create_connection(db_file):
    conn = None
    try:
        conn = sqlite3.connect(db_file)
    except Error as e:
        print(e)

    return conn

if __name__ == '__main__':
    app.run(host='0.0.0.0', port='8000', debug=True)


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
from absl import logging

import enum
import numpy as np
import six
from pysc2.lib import actions
from pysc2.lib import colors
from pysc2.lib import named_array
from pysc2.lib import point
from pysc2.lib import static_data
from pysc2.lib import stopwatch
from pysc2.lib import transform

from s2clientprotocol import raw_pb2 as sc_raw
from pysc2.lib import unit_controls
from s2clientprotocol import sc2api_pb2 as sc_pb

sw = stopwatch.sw

class FeatureType(enum.Enum):
  SCALAR = 1
  CATEGORICAL = 2

class PlayerRelative(enum.IntEnum):
  """The values for the `player_relative` feature layers."""
  NONE = 0
  SELF = 1
  ALLY = 2
  NEUTRAL = 3
  ENEMY = 4

class Visibility(enum.IntEnum):
  """Values for the `visibility` feature layers."""
  HIDDEN = 0
  SEEN = 1
  VISIBLE = 2

class Effects(enum.IntEnum):
  """Values for the `effects` feature layer."""
  # pylint: disable=invalid-name
  PsiStorm = 1
  GuardianShield = 2
  TemporalFieldGrowing = 3
  TemporalField = 4
  ThermalLance = 5
  ScannerSweep = 6
  NukeDot = 7
  LiberatorDefenderZoneSetup = 8
  LiberatorDefenderZone = 9
  BlindingCloud = 10
  CorrosiveBile = 11
  LurkerSpines = 12
  # pylint: enable=invalid-name

class ScoreCumulative(enum.IntEnum):
  """Indices into the `score_cumulative` observation."""
  score = 0
  idle_production_time = 1
  idle_worker_time = 2
  total_value_units = 3
  total_value_structures = 4
  killed_value_units = 5
  killed_value_structures = 6
  collected_minerals = 7
  collected_vespene = 8
  collection_rate_minerals = 9
  collection_rate_vespene = 10
  spent_minerals = 11
  spent_vespene = 12

class Player(enum.IntEnum):
  """Indices into the `player` observation."""
  player_id = 0
  minerals = 1
  vespene = 2
  food_used = 3
  food_cap = 4
  food_army = 5
  food_workers = 6
  idle_worker_count = 7
  army_count = 8
  warp_gate_count = 9
  larva_count = 10

class UnitLayer(enum.IntEnum):
  """Indices into the unit layers in the observations."""
  unit_type = 0
  player_relative = 1
  health = 2
  shields = 3
  energy = 4
  transport_slots_taken = 5
  build_progress = 6

class FeatureUnit(enum.IntEnum):
  """Indices for the `feature_unit` observations."""
  unit_type = 0
  alliance = 1
  health = 2
  shield = 3
  energy = 4
  cargo_space_taken = 5
  build_progress = 6
  health_ratio = 7
  shield_ratio = 8
  energy_ratio = 9
  display_type = 10
  owner = 11
  x = 12
  y = 13
  facing = 14
  radius = 15
  cloak = 16
  is_selected = 17
  is_blip = 18
  is_powered = 19
  mineral_contents = 20
  vespene_contents = 21
  cargo_space_max = 22
  assigned_harvesters = 23
  ideal_harvesters = 24
  weapon_cooldown = 25

class Feature(collections.namedtuple(
    "Feature", ["index", "name", "layer_set", "full_name", "scale", "type",
                "palette", "clip"])):
  """Define properties of a feature layer.

  Attributes:
    index: Index of this layer into the set of layers.
    name: The name of the layer within the set.
    layer_set: Which set of feature layers to look at in the observation proto.
    full_name: The full name including for visualization.
    scale: Max value (+1) of this layer, used to scale the values.
    type: A FeatureType for scalar vs categorical.
    palette: A color palette for rendering.
    clip: Whether to clip the values for coloring.
  """
  __slots__ = ()

  dtypes = {
      1: np.uint8,
      8: np.uint8,
      16: np.uint16,
      32: np.int32,
  }

  def unpack(self, obs):
    """Return a correctly shaped numpy array for this feature."""
    planes = getattr(obs.feature_layer_data, self.layer_set)
    plane = getattr(planes, self.name)
    return self.unpack_layer(plane)

  @staticmethod
  @sw.decorate
  def unpack_layer(plane):
    """Return a correctly shaped numpy array given the feature layer bytes."""
    size = point.Point.build(plane.size)
    if size == (0, 0):
      # New layer that isn't implemented in this SC2 version.
      return None
    data = np.frombuffer(plane.data, dtype=Feature.dtypes[plane.bits_per_pixel])
    if plane.bits_per_pixel == 1:
      data = np.unpackbits(data)
      if data.shape[0] != size.x * size.y:
        # This could happen if the correct length isn't a multiple of 8, leading
        # to some padding bits at the end of the string which are incorrectly
        # interpreted as data.
        data = data[:size.x * size.y]
    return data.reshape(size.y, size.x)

  @staticmethod
  @sw.decorate
  def unpack_rgb_image(plane):
    """Return a correctly shaped numpy array given the image bytes."""
    assert plane.bits_per_pixel == 24, "{} != 24".format(plane.bits_per_pixel)
    size = point.Point.build(plane.size)
    data = np.frombuffer(plane.data, dtype=np.uint8)
    return data.reshape(size.y, size.x, 3)

  @sw.decorate
  def color(self, plane):
    if self.clip:
      plane = np.clip(plane, 0, self.scale - 1)
    return self.palette[plane]

class ScreenFeatures(collections.namedtuple("ScreenFeatures", [
    "height_map", "visibility_map", "creep", "power", "player_id",
    "player_relative", "unit_type", "selected", "unit_hit_points",
    "unit_hit_points_ratio", "unit_energy", "unit_energy_ratio", "unit_shields",
    "unit_shields_ratio", "unit_density", "unit_density_aa", "effects"])):
  """The set of screen feature layers."""
  __slots__ = ()

  def __new__(cls, **kwargs):
    feats = {}
    for name, (scale, type_, palette, clip) in six.iteritems(kwargs):
      feats[name] = Feature(
          index=ScreenFeatures._fields.index(name),
          name=name,
          layer_set="renders",
          full_name="screen " + name,
          scale=scale,
          type=type_,
          palette=palette(scale) if callable(palette) else palette,
          clip=clip)
    return super(ScreenFeatures, cls).__new__(cls, **feats)  # pytype: disable=missing-parameter

class MinimapFeatures(collections.namedtuple("MinimapFeatures", [
    "height_map", "visibility_map", "creep", "camera", "player_id",
    "player_relative", "selected", "unit_type", "alerts", "pathable",
    "buildable"])):
  """The set of minimap feature layers."""
  __slots__ = ()

  def __new__(cls, **kwargs):
    feats = {}
    for name, (scale, type_, palette) in six.iteritems(kwargs):
      feats[name] = Feature(
          index=MinimapFeatures._fields.index(name),
          name=name,
          layer_set="minimap_renders",
          full_name="minimap " + name,
          scale=scale,
          type=type_,
          palette=palette(scale) if callable(palette) else palette,
          clip=False)
    return super(MinimapFeatures, cls).__new__(cls, **feats)  # pytype: disable=missing-parameter

SCREEN_FEATURES = ScreenFeatures(
    height_map=(256, FeatureType.SCALAR, colors.winter, False),
    visibility_map=(4, FeatureType.CATEGORICAL,
                    colors.VISIBILITY_PALETTE, False),
    creep=(2, FeatureType.CATEGORICAL, colors.CREEP_PALETTE, False),
    power=(2, FeatureType.CATEGORICAL, colors.POWER_PALETTE, False),
    player_id=(17, FeatureType.CATEGORICAL,
               colors.PLAYER_ABSOLUTE_PALETTE, False),
    player_relative=(5, FeatureType.CATEGORICAL,
                     colors.PLAYER_RELATIVE_PALETTE, False),
    unit_type=(max(static_data.UNIT_TYPES) + 1, FeatureType.CATEGORICAL,
               colors.unit_type, False),
    selected=(2, FeatureType.CATEGORICAL, colors.SELECTED_PALETTE, False),
    unit_hit_points=(1600, FeatureType.SCALAR, colors.hot, True),
    unit_hit_points_ratio=(256, FeatureType.SCALAR, colors.hot, False),
    unit_energy=(1000, FeatureType.SCALAR, colors.hot, True),
    unit_energy_ratio=(256, FeatureType.SCALAR, colors.hot, False),
    unit_shields=(1000, FeatureType.SCALAR, colors.hot, True),
    unit_shields_ratio=(256, FeatureType.SCALAR, colors.hot, False),
    unit_density=(16, FeatureType.SCALAR, colors.hot, True),
    unit_density_aa=(256, FeatureType.SCALAR, colors.hot, False),
    effects=(16, FeatureType.CATEGORICAL, colors.effects, False),
)

MINIMAP_FEATURES = MinimapFeatures(
    height_map=(256, FeatureType.SCALAR, colors.winter),
    visibility_map=(4, FeatureType.CATEGORICAL, colors.VISIBILITY_PALETTE),
    creep=(2, FeatureType.CATEGORICAL, colors.CREEP_PALETTE),
    camera=(2, FeatureType.CATEGORICAL, colors.CAMERA_PALETTE),
    player_id=(17, FeatureType.CATEGORICAL, colors.PLAYER_ABSOLUTE_PALETTE),
    player_relative=(5, FeatureType.CATEGORICAL,
                     colors.PLAYER_RELATIVE_PALETTE),
    selected=(2, FeatureType.CATEGORICAL, colors.winter),
    unit_type=(max(static_data.UNIT_TYPES) + 1, FeatureType.CATEGORICAL,
               colors.unit_type),
    alerts=(2, FeatureType.CATEGORICAL, colors.winter),
    pathable=(2, FeatureType.CATEGORICAL, colors.winter),
    buildable=(2, FeatureType.CATEGORICAL, colors.winter),
)

def _to_point(dims):
  """Convert (width, height) or size -> point.Point."""
  assert dims

  if isinstance(dims, (tuple, list)):
    if len(dims) != 2:
      raise ValueError(
          "A two element tuple or list is expected here, got {}.".format(dims))
    else:
      width = int(dims[0])
      height = int(dims[1])
      if width <= 0 or height <= 0:
        raise ValueError("Must specify +ve dims, got {}.".format(dims))
      else:
        return point.Point(width, height)
  else:
    size = int(dims)
    if size <= 0:
      raise ValueError(
          "Must specify a +ve value for size, got {}.".format(dims))
    else:
      return point.Point(size, size)

class Dimensions(object):
  """Screen and minimap dimensions configuration.

  Both screen and minimap must be specified. Sizes must be positive.
  Screen size must be greater than or equal to minimap size in both dimensions.

  Args:
    screen: A (width, height) int tuple or a single int to be used for both.
    minimap: A (width, height) int tuple or a single int to be used for both.
  """

  def __init__(self, screen=None, minimap=None):
    if not screen or not minimap:
      raise ValueError(
          "screen and minimap must both be set, screen={}, minimap={}".format(
              screen, minimap))

    self._screen = _to_point(screen)
    self._minimap = _to_point(minimap)

    if self._screen.x < self._minimap.x or self._screen.y < self._minimap.y:
      raise ValueError(
          "Screen (%s) can't be smaller than the minimap (%s)." % (
              self._screen, self._minimap))

  @property
  def screen(self):
    return self._screen

  @property
  def minimap(self):
    return self._minimap

  def __repr__(self):
    return "Dimensions(screen={}, minimap={})".format(self.screen, self.minimap)

class AgentInterfaceFormat(object):
  """Observation and action interface format specific to a particular agent."""

  def __init__(
      self,
      feature_dimensions=None,
      rgb_dimensions=None,
      action_space=None,
      camera_width_world_units=None,
      use_feature_units=False,
      hide_specific_actions=True):
    """Initializer.

    Args:
      feature_dimensions: Feature layer `Dimension`s. Either this or
          rgb_dimensions (or both) must be set.
      rgb_dimensions: RGB `Dimension`. Either this or feature_dimensions
          (or both) must be set.
      action_space: If you pass both feature and rgb sizes, then you must also
          specify which you want to use for your actions as an ActionSpace enum.
      camera_width_world_units: The width of your screen in world units. If your
          feature_dimensions.screen=(64, 48) and camera_width is 24, then each
          px represents 24 / 64 = 0.375 world units in each of x and y.
          It'll then represent a camera of size (24, 0.375 * 48) = (24, 18)
          world units.
      use_feature_units: Whether to include feature unit data in observations.
      hide_specific_actions: [bool] Some actions (eg cancel) have many
          specific versions (cancel this building, cancel that spell) and can
          be represented in a more general form. If a specific action is
          available, the general will also be available. If you set
          `hide_specific_actions` to False, the specific versions will also be
          available, but if it's True, the specific ones will be hidden.
          Similarly, when transforming back, a specific action will be returned
          as the general action. This simplifies the action space, though can
          lead to some actions in replays not being exactly representable using
          only the general actions.

    Raises:
      ValueError: if the parameters are inconsistent.
    """

    if not feature_dimensions and not rgb_dimensions:
      raise ValueError("Must set either the feature layer or rgb dimensions.")

    if action_space:
      if not isinstance(action_space, actions.ActionSpace):
        raise ValueError("action_space must be of type ActionSpace.")

      if ((action_space == actions.ActionSpace.FEATURES and
           not feature_dimensions) or
          (action_space == actions.ActionSpace.RGB and
           not rgb_dimensions)):
        raise ValueError(
            "Action space must match the observations, action space={}, "
            "feature_dimensions={}, rgb_dimensions={}".format(
                action_space, feature_dimensions, rgb_dimensions))
    else:
      if feature_dimensions and rgb_dimensions:
        raise ValueError(
            "You must specify the action space if you have both screen and "
            "rgb observations.")
      elif feature_dimensions:
        action_space = actions.ActionSpace.FEATURES
      else:
        action_space = actions.ActionSpace.RGB

    self._feature_dimensions = feature_dimensions
    self._rgb_dimensions = rgb_dimensions
    self._action_space = action_space
    self._camera_width_world_units = camera_width_world_units or 24
    self._use_feature_units = use_feature_units
    self._hide_specific_actions = hide_specific_actions

    if action_space == actions.ActionSpace.FEATURES:
      self._action_dimensions = feature_dimensions
    else:
      self._action_dimensions = rgb_dimensions

  @property
  def feature_dimensions(self):
    return self._feature_dimensions

  @property
  def rgb_dimensions(self):
    return self._rgb_dimensions

  @property
  def action_space(self):
    return self._action_space

  @property
  def camera_width_world_units(self):
    return self._camera_width_world_units

  @property
  def use_feature_units(self):
    return self._use_feature_units

  @property
  def hide_specific_actions(self):
    return self._hide_specific_actions

  @property
  def action_dimensions(self):
    return self._action_dimensions

def parse_agent_interface_format(
    feature_screen=None,
    feature_minimap=None,
    rgb_screen=None,
    rgb_minimap=None,
    action_space=None,
    camera_width_world_units=None,
    use_feature_units=False):
  """Creates an AgentInterfaceFormat object from keyword args.

  Convenient when using dictionaries or command-line arguments for config.

  Note that the feature_* and rgb_* properties define the respective spatial
  observation dimensions and accept:
      * None or 0 to disable that spatial observation.
      * A single int for a square observation with that side length.
      * A (int, int) tuple for a rectangular (width, height) observation.

  Args:
    feature_screen: If specified, so must feature_minimap be.
    feature_minimap: If specified, so must feature_screen be.
    rgb_screen: If specified, so must rgb_minimap be.
    rgb_minimap: If specified, so must rgb_screen be.
    action_space: ["FEATURES", "RGB"].
    camera_width_world_units: An int.
    use_feature_units: A boolean, defaults to False.

  Returns:
    An `AgentInterfaceFormat` object.

  Raises:
    ValueError: If an invalid parameter is specified.
  """
  if feature_screen or feature_minimap:
    feature_dimensions = Dimensions(
        screen=feature_screen,
        minimap=feature_minimap)
  else:
    feature_dimensions = None

  if rgb_screen or rgb_minimap:
    rgb_dimensions = Dimensions(
        screen=rgb_screen,
        minimap=rgb_minimap)
  else:
    rgb_dimensions = None

  return AgentInterfaceFormat(
      feature_dimensions=feature_dimensions,
      rgb_dimensions=rgb_dimensions,
      action_space=(action_space and actions.ActionSpace[action_space.upper()]),
      camera_width_world_units=camera_width_world_units,
      use_feature_units=use_feature_units
  )

def features_from_game_info(
    game_info,
    use_feature_units=False,
    action_space=None,
    hide_specific_actions=True):
  """Construct a Features object using data extracted from game info.

  Args:
    game_info: A `sc_pb.ResponseGameInfo` from the game.
    use_feature_units: Whether to include the feature unit observation.
    action_space: If you pass both feature and rgb sizes, then you must also
        specify which you want to use for your actions as an ActionSpace enum.
    hide_specific_actions: [bool] Some actions (eg cancel) have many
        specific versions (cancel this building, cancel that spell) and can
        be represented in a more general form. If a specific action is
        available, the general will also be available. If you set
        `hide_specific_actions` to False, the specific versions will also be
        available, but if it's True, the specific ones will be hidden.
        Similarly, when transforming back, a specific action will be returned
        as the general action. This simplifies the action space, though can
        lead to some actions in replays not being exactly representable using
        only the general actions.

  Returns:
    A features object matching the specified parameterisation.

  """

  if game_info.options.HasField("feature_layer"):
    fl_opts = game_info.options.feature_layer
    feature_dimensions = Dimensions(
        screen=(fl_opts.resolution.x, fl_opts.resolution.y),
        minimap=(fl_opts.minimap_resolution.x, fl_opts.minimap_resolution.y))
  else:
    feature_dimensions = None

  if game_info.options.HasField("render"):
    rgb_opts = game_info.options.render
    rgb_dimensions = Dimensions(
        screen=(rgb_opts.resolution.x, rgb_opts.resolution.y),
        minimap=(rgb_opts.minimap_resolution.x, rgb_opts.minimap_resolution.y))
  else:
    rgb_dimensions = None

  map_size = game_info.start_raw.map_size
  camera_width_world_units = game_info.options.feature_layer.width

  return Features(
      agent_interface_format=AgentInterfaceFormat(
          feature_dimensions=feature_dimensions,
          rgb_dimensions=rgb_dimensions,
          use_feature_units=use_feature_units,
          camera_width_world_units=camera_width_world_units,
          action_space=action_space,
          hide_specific_actions=hide_specific_actions),
      map_size=map_size)

def _init_valid_functions(action_dimensions):
  """Initialize ValidFunctions and set up the callbacks."""
  sizes = {
      "screen": tuple(int(i) for i in action_dimensions.screen),
      "screen2": tuple(int(i) for i in action_dimensions.screen),
      "minimap": tuple(int(i) for i in action_dimensions.minimap),
  }

  types = actions.Arguments(*[
      actions.ArgumentType.spec(t.id, t.name, sizes.get(t.name, t.sizes))
      for t in actions.TYPES])

  functions = actions.Functions([
      actions.Function.spec(f.id, f.name, tuple(types[t.id] for t in f.args))
      for f in actions.FUNCTIONS])

  return actions.ValidActions(types, functions)

class Features(object):
  """Render feature layers from SC2 Observation protos into numpy arrays.

  This has the implementation details of how to render a starcraft environment.
  It translates between agent action/observation formats and starcraft
  action/observation formats, which should not be seen by agent authors. The
  starcraft protos contain more information than they should have access to.

  This is outside of the environment so that it can also be used in other
  contexts, eg a supervised dataset pipeline.
  """

  def __init__(self, agent_interface_format=None, map_size=None):
    """Initialize a Features instance matching the specified interface format.

    Args:
      agent_interface_format: See the documentation for `AgentInterfaceFormat`.
      map_size: The size of the map in world units, needed for feature_units.

    Raises:
      ValueError: if agent_interface_format isn't specified.
      ValueError: if map_size isn't specified when use_feature_units is.
    """
    if not agent_interface_format:
      raise ValueError("Please specify agent_interface_format")

    self._agent_interface_format = agent_interface_format

    if agent_interface_format.use_feature_units:
      self._init_camera(
          agent_interface_format.feature_dimensions,
          map_size,
          agent_interface_format.camera_width_world_units)

    self._valid_functions = _init_valid_functions(
        agent_interface_format.action_dimensions)

  def _init_camera(
      self, feature_dimensions, map_size, camera_width_world_units):
    """Initialize the feature_units camera."""
    if not map_size or not camera_width_world_units:
      raise ValueError(
          "Either pass the game_info with raw enabled, or map_size and "
          "camera_width_world_units in order to use feature_units.")
    map_size = point.Point.build(map_size)
    self._world_to_world_tl = transform.Linear(point.Point(1, -1),
                                               point.Point(0, map_size.y))
    self._world_tl_to_world_camera_rel = transform.Linear(offset=-map_size / 4)
    world_camera_rel_to_feature_screen = transform.Linear(
        feature_dimensions.screen / camera_width_world_units,
        feature_dimensions.screen / 2)
    self._world_to_feature_screen_px = transform.Chain(
        self._world_to_world_tl,
        self._world_tl_to_world_camera_rel,
        world_camera_rel_to_feature_screen,
        transform.PixelToCoord())

  def _update_camera(self, camera_center):
    """Update the camera transform based on the new camera center."""
    self._world_tl_to_world_camera_rel.offset = (
        -self._world_to_world_tl.fwd_pt(camera_center) *
        self._world_tl_to_world_camera_rel.scale)

  def observation_spec(self):
    """The observation spec for the SC2 environment.

    It's worth noting that the image-like observations are in y,x/row,column
    order which is different than the actions which are in x,y order. This is
    due to conflicting conventions, and to facilitate printing of the images.

    Returns:
      The dict of observation names to their tensor shapes. Shapes with a 0 can
      vary in length, for example the number of valid actions depends on which
      units you have selected.
    """
    obs_spec = named_array.NamedDict({
        "action_result": (0,),  # See error.proto: ActionResult.
        "alerts": (0,),  # See sc2api.proto: Alert.
        "available_actions": (0,),
        "build_queue": (0, len(UnitLayer)),  # pytype: disable=wrong-arg-types
        "cargo": (0, len(UnitLayer)),  # pytype: disable=wrong-arg-types
        "cargo_slots_available": (1,),
        "control_groups": (10, 2),
        "game_loop": (1,),
        "last_actions": (0,),
        "multi_select": (0, len(UnitLayer)),  # pytype: disable=wrong-arg-types
        "player": (len(Player),),  # pytype: disable=wrong-arg-types
        "score_cumulative": (len(ScoreCumulative),),  # pytype: disable=wrong-arg-types
        "single_select": (0, len(UnitLayer)),  # Only (n, 7) for n in (0, 1).  # pytype: disable=wrong-arg-types
        # "units": [n],  units have an array of object Unit which define in unit_control.py
    })

    aif = self._agent_interface_format

    if aif.feature_dimensions:
      obs_spec["feature_screen"] = (len(SCREEN_FEATURES),
                                    aif.feature_dimensions.screen.y,
                                    aif.feature_dimensions.screen.x)

      obs_spec["feature_minimap"] = (len(MINIMAP_FEATURES),
                                     aif.feature_dimensions.minimap.y,
                                     aif.feature_dimensions.minimap.x)
    if aif.rgb_dimensions:
      obs_spec["rgb_screen"] = (aif.rgb_dimensions.screen.y,
                                aif.rgb_dimensions.screen.x,
                                3)
      obs_spec["rgb_minimap"] = (aif.rgb_dimensions.minimap.y,
                                 aif.rgb_dimensions.minimap.x,
                                 3)
    if aif.use_feature_units:
      obs_spec["feature_units"] = (0, len(FeatureUnit))  # pytype: disable=wrong-arg-types
    return obs_spec

  def action_spec(self):
    """The action space pretty complicated and fills the ValidFunctions."""
    return self._valid_functions

  @sw.decorate
  def transform_obs(self, obs):
    """Render some SC2 observations into something an agent can handle."""
    empty = np.array([], dtype=np.int32).reshape((0, 7))
    out = named_array.NamedDict({  # Fill out some that are sometimes empty.
        "single_select": empty,
        "multi_select": empty,
        "build_queue": empty,
        "cargo": empty,
        "cargo_slots_available": np.array([0], dtype=np.int32),
    })

    def or_zeros(layer, size):
      if layer is not None:
        return layer.astype(np.int32, copy=False)
      else:
        return np.zeros((size.y, size.x), dtype=np.int32)

    aif = self._agent_interface_format

    if aif.feature_dimensions:
      out["feature_screen"] = named_array.NamedNumpyArray(
          np.stack(or_zeros(f.unpack(obs.observation),
                            aif.feature_dimensions.screen)
                   for f in SCREEN_FEATURES),
          names=[ScreenFeatures, None, None])
      out["feature_minimap"] = named_array.NamedNumpyArray(
          np.stack(or_zeros(f.unpack(obs.observation),
                            aif.feature_dimensions.minimap)
                   for f in MINIMAP_FEATURES),
          names=[MinimapFeatures, None, None])

    if aif.rgb_dimensions:
      out["rgb_screen"] = Feature.unpack_rgb_image(
          obs.observation.render_data.map).astype(np.int32)
      out["rgb_minimap"] = Feature.unpack_rgb_image(
          obs.observation.render_data.minimap).astype(np.int32)

    out["last_actions"] = np.array(
        [self.reverse_action(a).function for a in obs.actions],
        dtype=np.int32)

    out["action_result"] = np.array([o.result for o in obs.action_errors],
                                    dtype=np.int32)

    out["alerts"] = np.array(obs.observation.alerts, dtype=np.int32)

    out["game_loop"] = np.array([obs.observation.game_loop], dtype=np.int32)

    score_details = obs.observation.score.score_details
    out["score_cumulative"] = named_array.NamedNumpyArray([
        obs.observation.score.score,
        score_details.idle_production_time,
        score_details.idle_worker_time,
        score_details.total_value_units,
        score_details.total_value_structures,
        score_details.killed_value_units,
        score_details.killed_value_structures,
        score_details.collected_minerals,
        score_details.collected_vespene,
        score_details.collection_rate_minerals,
        score_details.collection_rate_vespene,
        score_details.spent_minerals,
        score_details.spent_vespene,
    ], names=ScoreCumulative, dtype=np.int32)

    player = obs.observation.player_common
    out["player"] = named_array.NamedNumpyArray([
        player.player_id,
        player.minerals,
        player.vespene,
        player.food_used,
        player.food_cap,
        player.food_army,
        player.food_workers,
        player.idle_worker_count,
        player.army_count,
        player.warp_gate_count,
        player.larva_count,
    ], names=Player, dtype=np.int32)

    def unit_vec(u):
      return np.array((
          u.unit_type,
          u.player_relative,
          u.health,
          u.shields,
          u.energy,
          u.transport_slots_taken,
          int(u.build_progress * 100),  # discretize
      ), dtype=np.int32)

    ui = obs.observation.ui_data

    with sw("ui"):
      groups = np.zeros((10, 2), dtype=np.int32)
      for g in ui.groups:
        groups[g.control_group_index, :] = (g.leader_unit_type, g.count)
      out["control_groups"] = groups

      if ui.single:
        out["single_select"] = named_array.NamedNumpyArray(
            [unit_vec(ui.single.unit)], [None, UnitLayer])

      if ui.multi and ui.multi.units:
        out["multi_select"] = named_array.NamedNumpyArray(
            [unit_vec(u) for u in ui.multi.units], [None, UnitLayer])

      if ui.cargo and ui.cargo.passengers:
        out["single_select"] = named_array.NamedNumpyArray(
            [unit_vec(ui.single.unit)], [None, UnitLayer])
        out["cargo"] = named_array.NamedNumpyArray(
            [unit_vec(u) for u in ui.cargo.passengers], [None, UnitLayer])
        out["cargo_slots_available"] = np.array([ui.cargo.slots_available],
                                                dtype=np.int32)

      if ui.production and ui.production.build_queue:
        out["single_select"] = named_array.NamedNumpyArray(
            [unit_vec(ui.production.unit)], [None, UnitLayer])
        out["build_queue"] = named_array.NamedNumpyArray(
            [unit_vec(u) for u in ui.production.build_queue],
            [None, UnitLayer])

    def feature_unit_vec(u):
      screen_pos = self._world_to_feature_screen_px.fwd_pt(
          point.Point.build(u.pos))
      screen_radius = self._world_to_feature_screen_px.fwd_dist(u.radius)
      return np.array((
          # Match unit_vec order
          u.unit_type,
          u.alliance,  # Self = 1, Ally = 2, Neutral = 3, Enemy = 4
          u.health,
          u.shield,
          u.energy,
          u.cargo_space_taken,
          int(u.build_progress * 100),  # discretize

          # Resume API order
          int(u.health / u.health_max * 255) if u.health_max > 0 else 0,
          int(u.shield / u.shield_max * 255) if u.shield_max > 0 else 0,
          int(u.energy / u.energy_max * 255) if u.energy_max > 0 else 0,
          u.display_type,  # Visible = 1, Snapshot = 2, Hidden = 3
          u.owner,  # 1-15, 16 = neutral
          screen_pos.x,
          screen_pos.y,
          u.facing,
          screen_radius,
          u.cloak,  # Cloaked = 1, CloakedDetected = 2, NotCloaked = 3
          u.is_selected,
          u.is_blip,
          u.is_powered,
          u.mineral_contents,
          u.vespene_contents,

          # Not populated for enemies or neutral
          u.cargo_space_max,
          u.assigned_harvesters,
          u.ideal_harvesters,
          u.weapon_cooldown,
      ), dtype=np.int32)

    raw = obs.observation.raw_data

    if aif.use_feature_units:
      with sw("feature_units"):
        # Update the camera location so we can calculate world to screen pos
        self._update_camera(point.Point.build(raw.player.camera))
        feature_units = []
        for u in raw.units:
          if u.is_on_screen and u.display_type != sc_raw.Hidden:
            feature_units.append(feature_unit_vec(u))
        out["feature_units"] = named_array.NamedNumpyArray(
            feature_units, [None, FeatureUnit], dtype=np.int32)

    out["available_actions"] = np.array(self.available_actions(obs.observation),
                                        dtype=np.int32)
    out["units"] = self.transform_unit_control(raw)
    out["raw_data"] = raw
    return out

  @sw.decorate
  def transform_unit_control(self, raw):
    units = []
    for u in raw.units:
      units.append(unit_controls.Unit(u=u))
    return units

  @sw.decorate
  def available_actions(self, obs):
    """Return the list of available action ids."""
    available_actions = set()
    hide_specific_actions = self._agent_interface_format.hide_specific_actions
    for i, func in six.iteritems(actions.FUNCTIONS_AVAILABLE):
      if func.avail_fn(obs):
        available_actions.add(i)
    for a in obs.abilities:
      if a.ability_id not in actions.ABILITY_IDS:
        logging.warning("Unknown ability %s seen as available.", a.ability_id)
        continue
      for func in actions.ABILITY_IDS[a.ability_id]:
        if func.function_type in actions.POINT_REQUIRED_FUNCS[a.requires_point]:
          if func.general_id == 0 or not hide_specific_actions:
            available_actions.add(func.id)
          if func.general_id != 0:  # Always offer generic actions.
            for general_func in actions.ABILITY_IDS[func.general_id]:
              if general_func.function_type is func.function_type:
                # Only the right type. Don't want to expose the general action
                # to minimap if only the screen version is available.
                available_actions.add(general_func.id)
                break
    return list(available_actions)

  @sw.decorate
  def transform_action(self, obs, func_call, skip_available=False):
    """Tranform an agent-style action to one that SC2 can consume.

    Args:
      obs: a `sc_pb.Observation` from the previous frame.
      func_call: a `FunctionCall` to be turned into a `sc_pb.Action`.
      skip_available: If True, assume the action is available. This should only
          be used for testing or if you expect to make actions that weren't
          valid at the last observation.

    Returns:
      a corresponding `sc_pb.Action`.

    Raises:
      ValueError: if the action doesn't pass validation.
    """
    func_id = func_call.function
    try:
      func = actions.FUNCTIONS[func_id]
    except KeyError:
      raise ValueError("Invalid function id: %s." % func_id)

    # Available?
    if not (skip_available or func_id in self.available_actions(obs)):
      raise ValueError("Function %s/%s is currently not available" % (
          func_id, func.name))

    # Right number of args?
    if len(func_call.arguments) != len(func.args):
      raise ValueError(
          "Wrong number of arguments for function: %s, got: %s" % (
              func, func_call.arguments))

    # Args are valid?
    aif = self._agent_interface_format
    for t, arg in zip(func.args, func_call.arguments):
      if t.name in ("screen", "screen2"):
        sizes = aif.action_dimensions.screen
      elif t.name == "minimap":
        sizes = aif.action_dimensions.minimap
      else:
        sizes = t.sizes

      if len(sizes) != len(arg):
        raise ValueError(
            "Wrong number of values for argument of %s, got: %s" % (
                func, func_call.arguments))

      for s, a in zip(sizes, arg):
        if not 0 <= a < s:
          raise ValueError("Argument is out of range for %s, got: %s" % (
              func, func_call.arguments))

    # Convert them to python types.
    kwargs = {type_.name: type_.fn(a)
              for type_, a in zip(func.args, func_call.arguments)}

    # Call the right callback to get an SC2 action proto.
    sc2_action = sc_pb.Action()
    kwargs["action"] = sc2_action
    kwargs["action_space"] = aif.action_space
    if func.ability_id:
      kwargs["ability_id"] = func.ability_id
    actions.FUNCTIONS[func_id].function_type(**kwargs)
    return sc2_action

  @sw.decorate
  def reverse_action(self, action):
    """Transform an SC2-style action into an agent-style action.

    This should be the inverse of `transform_action`.

    Args:
      action: a `sc_pb.Action` to be transformed.

    Returns:
      A corresponding `actions.FunctionCall`.

    Raises:
      ValueError: if it doesn't know how to transform this action.
    """
    FUNCTIONS = actions.FUNCTIONS  # pylint: disable=invalid-name

    aif = self._agent_interface_format

    def func_call_ability(ability_id, cmd_type, *args):
      """Get the function id for a specific ability id and action type."""
      if ability_id not in actions.ABILITY_IDS:
        logging.warning("Unknown ability_id: %s. This is probably dance or "
                        "cheer, or some unknown new or map specific ability. "
                        "Treating it as a no-op.", ability_id)
        return FUNCTIONS.no_op()

      if aif.hide_specific_actions:
        general_id = next(iter(actions.ABILITY_IDS[ability_id])).general_id
        if general_id:
          ability_id = general_id

      for func in actions.ABILITY_IDS[ability_id]:
        if func.function_type is cmd_type:
          return FUNCTIONS[func.id](*args)
      raise ValueError("Unknown ability_id: %s, type: %s. Likely a bug." % (
          ability_id, cmd_type.__name__))

    if action.HasField("action_ui"):
      act_ui = action.action_ui
      if act_ui.HasField("multi_panel"):
        return FUNCTIONS.select_unit(act_ui.multi_panel.type - 1,
                                     act_ui.multi_panel.unit_index)
      if act_ui.HasField("control_group"):
        return FUNCTIONS.select_control_group(
            act_ui.control_group.action - 1,
            act_ui.control_group.control_group_index)
      if act_ui.HasField("select_idle_worker"):
        return FUNCTIONS.select_idle_worker(act_ui.select_idle_worker.type - 1)
      if act_ui.HasField("select_army"):
        return FUNCTIONS.select_army(act_ui.select_army.selection_add)
      if act_ui.HasField("select_warp_gates"):
        return FUNCTIONS.select_warp_gates(
            act_ui.select_warp_gates.selection_add)
      if act_ui.HasField("select_larva"):
        return FUNCTIONS.select_larva()
      if act_ui.HasField("cargo_panel"):
        return FUNCTIONS.unload(act_ui.cargo_panel.unit_index)
      if act_ui.HasField("production_panel"):
        return FUNCTIONS.build_queue(act_ui.production_panel.unit_index)
      if act_ui.HasField("toggle_autocast"):
        return func_call_ability(act_ui.toggle_autocast.ability_id,
                                 actions.autocast)

    if (action.HasField("action_feature_layer") or
        action.HasField("action_render")):
      act_sp = actions.spatial(action, aif.action_space)
      if act_sp.HasField("camera_move"):
        coord = point.Point.build(act_sp.camera_move.center_minimap)
        return FUNCTIONS.move_camera(coord)
      if act_sp.HasField("unit_selection_point"):
        select_point = act_sp.unit_selection_point
        coord = point.Point.build(select_point.selection_screen_coord)
        return FUNCTIONS.select_point(select_point.type - 1, coord)
      if act_sp.HasField("unit_selection_rect"):
        select_rect = act_sp.unit_selection_rect
        # TODO(tewalds): After looking at some replays we should decide if
        # this is good enough. Maybe we need to simulate multiple actions or
        # merge the selection rects into a bigger one.
        tl = point.Point.build(select_rect.selection_screen_coord[0].p0)
        br = point.Point.build(select_rect.selection_screen_coord[0].p1)
        return FUNCTIONS.select_rect(select_rect.selection_add, tl, br)
      if act_sp.HasField("unit_command"):
        cmd = act_sp.unit_command
        queue = int(cmd.queue_command)
        if cmd.HasField("target_screen_coord"):
          coord = point.Point.build(cmd.target_screen_coord)
          return func_call_ability(cmd.ability_id, actions.cmd_screen,
                                   queue, coord)
        elif cmd.HasField("target_minimap_coord"):
          coord = point.Point.build(cmd.target_minimap_coord)
          return func_call_ability(cmd.ability_id, actions.cmd_minimap,
                                   queue, coord)
        else:
          return func_call_ability(cmd.ability_id, actions.cmd_quick, queue)

    if action.HasField("action_raw") or action.HasField("action_render"):
      raise ValueError("Unknown action:\n%s" % action)

    return FUNCTIONS.no_op()


from dgl.nn.functional import edge_softmax
from dgl import DGLError
import dgl.function as fn
import torch
import torch.nn as nn
import torch.nn.functional as F

class GATLayer(nn.Module):
    """Feature computation h => h'.

    **Linear transformation**:
        :math:`z_i=Wh_i`
    **Edge attention**:
        :math:`e_{ij}=LeakyReLU\ (a^T\cdot(z_i\ ||\ z_j\ ))`
    **Aggregation**:
        :math:`a_{ij}=\exp(e_{ij}\ )/(\sum_{k\in \mathcal{N}(i\ )}\exp(e_{ik}\ ))`

        :math:`h_i'=Relu\ (\sum_{j\in \mathcal{N}(i\ )} {a_{ij}\ z_j\ })`
    """
    def __init__(self, in_dim: int, out_dim: int,
                 feat_drop: float = 0., attn_drop: float = 0.,
                 negative_slope: float = 0.2):
        super(GATLayer, self).__init__()
        self.linear = nn.Linear(in_dim, out_dim, bias=False)
        self.attn_s = nn.Linear(out_dim, 1, bias=False)
        self.attn_d = nn.Linear(out_dim, 1, bias=False)
        self.feat_drop = nn.Dropout(feat_drop)
        self.attn_drop = nn.Dropout(attn_drop)
        self.leaky_relu = nn.LeakyReLU(negative_slope)
        self.reset_parameters()

    def reset_parameters(self):
        gain = nn.init.calculate_gain('relu')
        nn.init.xavier_normal_(self.linear.weight, gain=gain)
        nn.init.xavier_normal_(self.attn_s.weight, gain=gain)
        nn.init.xavier_normal_(self.attn_d.weight, gain=gain)

    def forward(self, graph, h):
        with graph.local_scope():
            h = self.feat_drop(h)
            z_s = self.linear(h)
            if graph.is_block:
                z_d = z_s[:graph.number_of_dst_nodes()]
            else:
                z_d = z_s

            # compute edge attention
            es = self.attn_s(z_s)
            ed = self.attn_d(z_d)
            graph.srcdata.update({'z': z_s, 'es': es})
            graph.dstdata.update({'ed': ed})
            graph.apply_edges(fn.u_add_v('es', 'ed', 'e'))
            e = self.leaky_relu(graph.edata.pop('e'))
            graph.edata['a'] = self.attn_drop(edge_softmax(graph, e))

            # message passing
            graph.update_all(fn.u_mul_e('z', 'a', 'm'), fn.sum('m', 'z'))

            return graph.dstdata['z']

class MultiHeadGATLayer(nn.Module):
    def __init__(self,
                 in_dim: int,
                 out_dim: int,
                 num_heads: int,
                 dropout: float = 0.2):
        super(MultiHeadGATLayer, self).__init__()
        self.heads = nn.ModuleList()
        for i in range(num_heads):
            self.heads.append(GATLayer(
                in_dim, out_dim, feat_drop=dropout, attn_drop=0.2))

    def forward(self, graph, h):
        batch = [head(graph, h) for head in self.heads]
        return torch.cat(batch, dim=-1)

class JKNLayer(nn.Module):
    """**Layer-level aggregation**
        :math:`\mathbf{h}_{final}=\phi(\mathbf{h}^1, \mathbf{h}^2, ...\ , \mathbf{h}^L\ )`

    * **Concatenation**: :math:`\phi = ||`

    * **Gated Recurrent Unit**:

        :math:`(\mathbf{f}^l,\mathbf{b}^l) = bidirectional\ GRU\ (\mathbf{h}^l\ )`

        Attention score for layers
        :math:`s^l = a^T\cdot(\mathbf{f}^l\ ||\ \mathbf{b}^l)`

        :math:`\mathbf{h}_{final} = softmax(\mathbf{s})\cdot(\mathbf{f}\ ||\ \mathbf{b})`

    Parameters
    ----------
    in_dim : int
        input feature size.
    out_dim : int
        output feature size.
    n_layers : int
        number of layers in multi-layer skip connection.
    aggregator_type : str
        the aggregation method. It supports "cat" and "gru". The default value is "cat".
    """
    def __init__(self, in_dim: int, out_dim: int,
                 n_layers: int,
                 aggregator_type='cat'):
        super(JKNLayer, self).__init__()
        if aggregator_type not in ['cat', 'gru']:
            raise KeyError('Aggregator type {} not supported.'.format(aggregator_type))
        self.__aggr_type = aggregator_type
        if aggregator_type == 'cat':
            self.linear = nn.Linear(in_dim * n_layers, out_dim, bias=False)
        if aggregator_type == 'gru':
            self.gru = nn.GRU(input_size=in_dim,
                              hidden_size=in_dim,
                              bidirectional=True)
            self.s = nn.Linear(in_dim * 2, 1, bias=False)
            self.linear = nn.Linear(in_dim * 2, out_dim, bias=False)

    def reset_parameters(self):
        gain = nn.init.calculate_gain('relu')
        nn.init.xavier_normal_(self.linear.weight, gain=gain)
        if self.__aggr_type == 'gru':
            nn.init.xavier_normal_(self.s.weight, gain=gain)

    def forward(self, h_list):
        if self.__aggr_type == 'cat':
            h_cat = torch.cat(h_list, dim=-1)
            return self.linear(h_cat)
        else:
            # gru
            h_seq, _ = self.gru(torch.stack(h_list, dim=0))
            w = self.s(h_seq)
            w = F.softmax(w, dim=-1)
            h_N = torch.einsum('lni,lnf->nf', w, h_seq)
            return self.linear(h_N)

class Model(nn.Module):
    def __init__(self, in_dim: int, h_dim: int, out_dim: int,
                 n_layers: int = 2, n_heads: int = 2,
                 dropout: float = 0.2, skip_connection='gru'):
        super(Model, self).__init__()
        self.input_layer = MultiHeadGATLayer(in_dim, h_dim, n_heads, dropout)
        if n_layers < 1:
            raise DGLError('Number of layers should be positive integer.')
        hid_dim = h_dim * n_heads
        n_h_layers = n_layers - 1
        self.layers = nn.ModuleList()
        for i in range(n_h_layers):
            self.layers.append(MultiHeadGATLayer(hid_dim, h_dim, n_heads, dropout))
        self.output_layer = JKNLayer(hid_dim, out_dim, n_layers, skip_connection)

    def forward(self, graph, in_feat):
        h = self.input_layer(graph, in_feat)
        h_list = [F.relu(h)]
        for layer in self.layers:
            h = layer(graph, F.relu(h))
            h_list.append(F.relu(h))
        return self.output_layer(h_list)


import numpy as np
import torch
from torch.autograd import Variable
import matplotlib.pyplot as plt

import copy
import random
import logging
import time
import datetime

from tqdm import tqdm
from games.gridworld import Gridworld as Game
from gridrl.utils import count_parameters

LOGGER = logging.getLogger(__name__)

class RLTrainer(object):

    def __init__(self, game_parameters, train_parameters):
        self.device = torch.device(
            'cuda:0' if torch.cuda.is_available() else 'cpu'
        )
        LOGGER.info('Device = {}'.format(self.device))
        self.action_set = Game.action_set
        self.game_size = game_parameters['size']
        self.game_mode = game_parameters['mode']
        self._configure(train_parameters)

    def _configure(self, train_d):
        # set these properties directly
        self.model = None
        self.target_model = None
        self.loss_fn = torch.nn.MSELoss(reduction='elementwise_mean')
        self.optimizer = None
        self.start_epoch = None
        self.start_step = None
        # train configuration
        self.batch_size = train_d.get('batch_size', 100)
        self.buffer = train_d.get('buffer', 500)
        self.ckpt_path = train_d.get('ckpt_path', './ckpt_deepq_targrep.tar')
        self.epsilon = train_d.get('epsilon', 1.0)
        self.gamma = train_d.get('gamma', 0.95)
        self.learning_rate = train_d.get('learning_rate', 1e-3)
        self.max_moves = train_d.get('max_moves', 50)
        self.min_epsilon = train_d.get('min_epsilon', 0.1)
        self.saved_losses_path = train_d.get('saved_losses_path', 'losses.npy')
        self.saved_winpct_path = train_d.get('saved_winpct_path', 'winpct.npy')
        self.target_network_update = train_d.get(
            'target_network_update', 500
        )
        self.show_progress = train_d.get('show_progress', False)
        self._f = tqdm if self.show_progress else (lambda x: x)
        # for visualizing training progress
        self.losses = []
        self.winpct = []
        try:
            losses_file = self.saved_losses_path
            winpct_file = self.saved_winpct_path
            losses_arr = np.load(losses_file)
            winpct_arr = np.load(winpct_file)
            self.losses.extend(list(map(tuple, losses_arr)))
            self.winpct.extend(list(map(tuple, winpct_arr)))
        except FileNotFoundError:
            pass

    def build_or_restore_model_and_optimizer(
        self, build_model_function, conv, epsilon=None
    ):
        '''
        params:
        * build_model_function - fn that returns the model to be used
        * conv - T/F - how to reshape game input
        '''
        self.model = build_model_function()
        LOGGER.info('model has {} parameters'.format(
            count_parameters(self.model)
        ))
        self.conv = conv
        self.optimizer = torch.optim.Adam(self.model.parameters(),
                                          lr=self.learning_rate)
        self.start_epoch = 0
        self.start_step = 0

        try:
            checkpoint = torch.load(self.ckpt_path)
            self.model.load_state_dict(checkpoint['model_state_dict'])
            self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
            self.start_epoch = checkpoint['epoch'] + 1
            self.start_step = checkpoint['step'] + 1
            self.epsilon = checkpoint['epsilon']
            LOGGER.info('Loaded checkpoint from {}'.format(self.ckpt_path))
        except FileNotFoundError:
            LOGGER.info('No checkpoint found...')

        if epsilon is not None:
            self.epsilon = epsilon

        LOGGER.info('Model state dict:')
        for param_tensor in self.model.state_dict():
            LOGGER.info(str(param_tensor) + '\t'
                        + str(self.model.state_dict()[param_tensor].size()))
        LOGGER.info('Optimizer state dict:')
        for var_name in self.optimizer.state_dict():
            LOGGER.info(str(var_name) + '\t'
                        + str(self.optimizer.state_dict()[var_name]))

        self.model.to(self.device)
        self.target_model = copy.deepcopy(self.model)
        self.target_model.to(self.device)

    def _new_game(self):
        game = Game(size=self.game_size, mode=self.game_mode)
        return game

    def _game_move_for_reward(self, game, action):
        game.makeMove(action)
        reward = game.reward()
        return reward

    def _check_continue(self, reward, move_count):
        if reward != Game.default_reward or move_count > self.max_moves:
            return 0, 0
        return move_count, 1

    def _get_torch_state(self, game, noise_factor=100.0):
        game_boardsz = game.board.size
        if self.conv:
            state_ = game.board.render_np() + \
                np.random.rand(4, game_boardsz, game_boardsz) / noise_factor
            state_ = state_.reshape(1, 4, game_boardsz, game_boardsz)
        else:
            full_game_size = 4 * game_boardsz * game_boardsz
            state_ = game.board.render_np().reshape(1, full_game_size) + \
                np.random.rand(1, full_game_size) / noise_factor
        state = Variable(torch.from_numpy(state_).float()).to(self.device)
        return state

    def _compute_action(self, qval):
        qval_ = qval.cpu().data.numpy()
        if np.random.rand() < self.epsilon:
            action_ = np.random.randint(0, 4)
        else:
            action_ = np.argmax(qval_)
        action_str = self.action_set[action_]
        return action_str, action_

    def _build_trainbatch(self, replay):
        minibatch = random.sample(replay, self.batch_size)
        X_train = Variable(torch.empty(self.batch_size, 4, dtype=torch.float))
        y_train = Variable(torch.empty(self.batch_size, 4, dtype=torch.float))
        # Fill X_train and y_train minibatch tensors by index `h` by
        # looping through memory and computing the Q-values before (X)
        # and after (y) each move
        h = 0
        for memory in minibatch:
            old_state, action_m, reward_m, new_state_m = memory
            old_qval = self.model(old_state)
            new_qval = self.target_model(new_state_m).cpu().data.numpy()
            max_qval = np.max(new_qval)
            y = torch.zeros((1, 4))
            y[:] = old_qval[:]
            if reward_m == -1:
                update = (reward_m + (self.gamma * max_qval))
            else:
                update = reward_m
            y[0][action_m] = update
            X_train[h] = old_qval
            y_train[h] = Variable(y)
            h += 1
        X_train, y_train = X_train.to(self.device), y_train.to(self.device)
        return X_train, y_train

    def _anneal_epsilon(self, epochs):
        if self.epsilon > self.min_epsilon:
            self.epsilon -= (1. / epochs)

    def _save_model(self, epoch, step):
        torch.save({
            'epoch': epoch,
            'model_state_dict': self.model.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'step': step,
            'epsilon': self.epsilon,
        }, self.ckpt_path)

    def _play_single_game(self):
        LOGGER.debug('Tesing single game:')
        i = 0
        game = self._new_game()
        state = self._get_torch_state(game)
        LOGGER.debug(' Initial state:')
        LOGGER.debug(str(game.display()))
        status = 1
        while status == 1:
            qval = self.model(state)
            action, action_ = self._compute_action(qval)
            LOGGER.debug('  Move #: %s; Taking action: %s' % (i, action))
            reward = self._game_move_for_reward(game, action)
            state = self._get_torch_state(game)
            LOGGER.debug(str(game.display()))
            if reward != Game.default_reward:
                status = 0
                if reward == Game.win_reward:
                    return 1
                elif reward == Game.loss_reward:
                    return 0
                else:
                    raise ValueError('Bad game reward.')
            i += 1
            if (i > self.max_moves):
                return 0

    def play_model(self, max_games=1000, step_count=None):
        wins = 0
        for i in self._f(range(max_games)):
            win = self._play_single_game()
            if win:
                wins += 1
        win_perc = float(wins) / max_games
        fs = '  step={:010d}, games={:04d}/{:04d}, win pct={:1.5f}'
        LOGGER.info(fs.format(step_count, wins, max_games, win_perc))
        if step_count is not None:
            self.winpct.append((step_count, win_perc))

    def train_model_with_target_replay(self, epochs):
        LOGGER.info('Running training...')
        LOGGER.info(' Start epoch = {}; start step = {}; epsilon = {}'.format(
            self.start_epoch, self.start_step, self.epsilon
        ))
        replay, step, c_step = [], self.start_step, 0
        t0 = time.perf_counter()
        for epoch in self._f(
            range(self.start_epoch, self.start_epoch + epochs)
        ):
            game = self._new_game()
            state = self._get_torch_state(game)
            status, move_count = 1, 0
            while status == 1:
                move_count, step, c_step = move_count + 1, step + 1, c_step + 1
                if c_step > self.target_network_update:
                    LOGGER.info(" ...updating target model")
                    self.target_model.load_state_dict(self.model.state_dict())
                    c_step = 0
                # Compute Q-value and use to make move decisions
                qval = self.model(state)
                action, action_ = self._compute_action(qval)
                reward = self._game_move_for_reward(game, action)
                new_state = self._get_torch_state(game)

                # save move information into replay buffer
                if len(replay) < self.buffer:
                    replay.append((state, action_, reward, new_state))
                else:
                    replay.pop(0)
                    replay.append((state, action_, reward, new_state))
                    X_train, y_train = self._build_trainbatch(replay)
                    loss = self.loss_fn(X_train, y_train)
                    self.optimizer.zero_grad()
                    loss.backward()
                    self.losses.append((step, loss.item()))
                    self.optimizer.step()
                    LOGGER.debug(
                        'epoch={:06d}, step={:08d}, loss={:04.8f}'.format(
                            epoch, step, loss.item()
                        ))

                state = new_state
                move_count, status = self._check_continue(reward, move_count)

            # after game/epoch, update epsilon, ...
            LOGGER.info('replay buffer length at epoch {} end = {}'.format(
                epoch, len(replay)
            ))
            self._anneal_epsilon(epochs)
            if epoch % 50 == 0:
                LOGGER.info(' Playing model at epoch {}'.format(epoch))
                self.play_model(step_count=step)
                self._save_model(epoch, step)

        # Final save
        self.play_model(step_count=step)
        self._save_model(epoch, step)
        t1 = time.perf_counter()
        LOGGER.info('Total train time for {} epochs = {:04.3f}s = {}'.format(
            epochs, t1 - t0, str(datetime.timedelta(seconds=t1-t0))
        ))

    def save_losses_and_winpct_plots(self, make_plot):

        def running_mean(x, N=500):
            cumsum = np.cumsum(np.insert(x, 0, 0))
            return (cumsum[N:] - cumsum[:-N]) / float(N)

        wins_arr = np.asarray(self.winpct)
        loss_arr = np.asarray(self.losses)
        np.save(self.saved_losses_path, loss_arr)
        np.save(self.saved_winpct_path, wins_arr)

        if make_plot:

            try:
                wins_stps = wins_arr[:, 0]
                wins_wins = wins_arr[:, 1]
                loss_stps = loss_arr[:, 0]
                loss_loss = loss_arr[:, 1]
                nrunning_mean = 5
                running_loss = running_mean(loss_loss, N=nrunning_mean)

                fig = plt.figure()
                gs = plt.GridSpec(1, 2)

                ax1 = plt.subplot(gs[0])
                ax1.scatter(wins_stps, wins_wins)
                ax1.set_xlabel('Steps')
                ax1.set_ylabel('Win percentage')
                ax1.set_ylim(0., 1.)
                ax2 = plt.subplot(gs[1])
                ax2.scatter(loss_stps[nrunning_mean - 1:], running_loss)
                ax2.set_xlabel('Steps')
                ax2.set_ylabel('Running Loss')
                ax2.set_ylim(0., 3.)

                fig.tight_layout()
                figname = 'deepq_targrep_%d.pdf' % (time.time())
                plt.savefig(figname, bbox_inches='tight')
            except IndexError as e:
                LOGGER.error(e)
                LOGGER.error('Empty loss/winpct arrays...')