import logging
from typing import List, Optional, Tuple
import torch
from fvcore.nn import sigmoid_focal_loss_jit
from torch import nn
from torch.nn import functional as F

from detectron2.layers import ShapeSpec, batched_nms
from detectron2.structures import Boxes, ImageList, Instances, pairwise_point_box_distance
from detectron2.utils.events import get_event_storage

from ..anchor_generator import DefaultAnchorGenerator
from ..backbone import Backbone
from ..box_regression import Box2BoxTransformLinear, _dense_box_regression_loss
from .dense_detector import DenseDetector
from .retinanet import RetinaNetHead

__all__ = ["FCOS"]

logger = logging.getLogger(__name__)

class FCOS(DenseDetector):
    """
    Implement FCOS in :paper:`fcos`.
    """

    def __init__(
        self,
        *,
        backbone: Backbone,
        head: nn.Module,
        head_in_features: Optional[List[str]] = None,
        box2box_transform=None,
        num_classes,
        center_sampling_radius: float = 1.5,
        focal_loss_alpha=0.25,
        focal_loss_gamma=2.0,
        test_score_thresh=0.2,
        test_topk_candidates=1000,
        test_nms_thresh=0.6,
        max_detections_per_image=100,
        pixel_mean,
        pixel_std,
    ):
        """
        Args:
            center_sampling_radius: radius of the "center" of a groundtruth box,
                within which all anchor points are labeled positive.
            Other arguments mean the same as in :class:`RetinaNet`.
        """
        super().__init__(
            backbone, head, head_in_features, pixel_mean=pixel_mean, pixel_std=pixel_std
        )

        self.num_classes = num_classes

        # FCOS uses one anchor point per location.
        # We represent the anchor point by a box whose size equals the anchor stride.
        feature_shapes = backbone.output_shape()
        fpn_strides = [feature_shapes[k].stride for k in self.head_in_features]
        self.anchor_generator = DefaultAnchorGenerator(
            sizes=[[k] for k in fpn_strides], aspect_ratios=[1.0], strides=fpn_strides
        )

        # FCOS parameterizes box regression by a linear transform,
        # where predictions are normalized by anchor stride (equal to anchor size).
        if box2box_transform is None:
            box2box_transform = Box2BoxTransformLinear(normalize_by_size=True)
        self.box2box_transform = box2box_transform

        self.center_sampling_radius = float(center_sampling_radius)

        # Loss parameters:
        self.focal_loss_alpha = focal_loss_alpha
        self.focal_loss_gamma = focal_loss_gamma

        # Inference parameters:
        self.test_score_thresh = test_score_thresh
        self.test_topk_candidates = test_topk_candidates
        self.test_nms_thresh = test_nms_thresh
        self.max_detections_per_image = max_detections_per_image

    def forward_training(self, images, features, predictions, gt_instances):
        # Transpose the Hi*Wi*A dimension to the middle:
        pred_logits, pred_anchor_deltas, pred_centerness = self._transpose_dense_predictions(
            predictions, [self.num_classes, 4, 1]
        )
        anchors = self.anchor_generator(features)
        gt_labels, gt_boxes = self.label_anchors(anchors, gt_instances)
        return self.losses(
            anchors, pred_logits, gt_labels, pred_anchor_deltas, gt_boxes, pred_centerness
        )

    @torch.no_grad()
    def _match_anchors(self, gt_boxes: Boxes, anchors: List[Boxes]):
        """
        Match ground-truth boxes to a set of multi-level anchors.

        Args:
            gt_boxes: Ground-truth boxes from instances of an image.
            anchors: List of anchors for each feature map (of different scales).

        Returns:
            torch.Tensor
                A tensor of shape `(M, R)`, given `M` ground-truth boxes and total
                `R` anchor points from all feature levels, indicating the quality
                of match between m-th box and r-th anchor. Higher value indicates
                better match.
        """
        # Naming convention: (M = ground-truth boxes, R = anchor points)
        # Anchor points are represented as square boxes of size = stride.
        num_anchors_per_level = [len(x) for x in anchors]
        anchors = Boxes.cat(anchors)  # (R, 4)
        anchor_centers = anchors.get_centers()  # (R, 2)
        anchor_sizes = anchors.tensor[:, 2] - anchors.tensor[:, 0]  # (R, )

        lower_bound = anchor_sizes * 4
        lower_bound[: num_anchors_per_level[0]] = 0
        upper_bound = anchor_sizes * 8
        upper_bound[-num_anchors_per_level[-1] :] = float("inf")

        gt_centers = gt_boxes.get_centers()

        # FCOS with center sampling: anchor point must be close enough to
        # ground-truth box center.
        center_dists = (anchor_centers[None, :, :] - gt_centers[:, None, :]).abs_()
        sampling_regions = self.center_sampling_radius * anchor_sizes[None, :]

        match_quality_matrix = center_dists.max(dim=2).values < sampling_regions

        pairwise_dist = pairwise_point_box_distance(anchor_centers, gt_boxes)
        pairwise_dist = pairwise_dist.permute(1, 0, 2)  # (M, R, 4)

        # The original FCOS anchor matching rule: anchor point must be inside GT.
        match_quality_matrix &= pairwise_dist.min(dim=2).values > 0

        # Multilevel anchor matching in FCOS: each anchor is only responsible
        # for certain scale range.
        pairwise_dist = pairwise_dist.max(dim=2).values
        match_quality_matrix &= (pairwise_dist > lower_bound[None, :]) & (
            pairwise_dist < upper_bound[None, :]
        )
        # Match the GT box with minimum area, if there are multiple GT matches.
        gt_areas = gt_boxes.area()  # (M, )

        match_quality_matrix = match_quality_matrix.to(torch.float32)
        match_quality_matrix *= 1e8 - gt_areas[:, None]
        return match_quality_matrix  # (M, R)

    @torch.no_grad()
    def label_anchors(self, anchors: List[Boxes], gt_instances: List[Instances]):
        """
        Same interface as :meth:`RetinaNet.label_anchors`, but implemented with FCOS
        anchor matching rule.

        Unlike RetinaNet, there are no ignored anchors.
        """

        gt_labels, matched_gt_boxes = [], []

        for inst in gt_instances:
            if len(inst) > 0:
                match_quality_matrix = self._match_anchors(inst.gt_boxes, anchors)

                # Find matched ground-truth box per anchor. Un-matched anchors are
                # assigned -1. This is equivalent to using an anchor matcher as used
                # in R-CNN/RetinaNet: `Matcher(thresholds=[1e-5], labels=[0, 1])`
                match_quality, matched_idxs = match_quality_matrix.max(dim=0)
                matched_idxs[match_quality < 1e-5] = -1

                matched_gt_boxes_i = inst.gt_boxes.tensor[matched_idxs.clip(min=0)]
                gt_labels_i = inst.gt_classes[matched_idxs.clip(min=0)]

                # Anchors with matched_idxs = -1 are labeled background.
                gt_labels_i[matched_idxs < 0] = self.num_classes
            else:
                matched_gt_boxes_i = torch.zeros_like(Boxes.cat(anchors).tensor)
                gt_labels_i = torch.full(
                    (len(matched_gt_boxes_i),),
                    fill_value=self.num_classes,
                    dtype=torch.long,
                    device=matched_gt_boxes_i.device,
                )

            gt_labels.append(gt_labels_i)
            matched_gt_boxes.append(matched_gt_boxes_i)

        return gt_labels, matched_gt_boxes

    def losses(
        self, anchors, pred_logits, gt_labels, pred_anchor_deltas, gt_boxes, pred_centerness
    ):
        """
        This method is almost identical to :meth:`RetinaNet.losses`, with an extra
        "loss_centerness" in the returned dict.
        """
        num_images = len(gt_labels)
        gt_labels = torch.stack(gt_labels)  # (M, R)

        pos_mask = (gt_labels >= 0) & (gt_labels != self.num_classes)
        num_pos_anchors = pos_mask.sum().item()
        get_event_storage().put_scalar("num_pos_anchors", num_pos_anchors / num_images)
        normalizer = self._ema_update("loss_normalizer", max(num_pos_anchors, 1), 300)

        # classification and regression loss
        gt_labels_target = F.one_hot(gt_labels, num_classes=self.num_classes + 1)[
            :, :, :-1
        ]  # no loss for the last (background) class
        loss_cls = sigmoid_focal_loss_jit(
            torch.cat(pred_logits, dim=1),
            gt_labels_target.to(pred_logits[0].dtype),
            alpha=self.focal_loss_alpha,
            gamma=self.focal_loss_gamma,
            reduction="sum",
        )

        loss_box_reg = _dense_box_regression_loss(
            anchors,
            self.box2box_transform,
            pred_anchor_deltas,
            gt_boxes,
            pos_mask,
            box_reg_loss_type="giou",
        )

        ctrness_targets = self.compute_ctrness_targets(anchors, gt_boxes)  # (M, R)
        pred_centerness = torch.cat(pred_centerness, dim=1).squeeze(dim=2)  # (M, R)
        ctrness_loss = F.binary_cross_entropy_with_logits(
            pred_centerness[pos_mask], ctrness_targets[pos_mask], reduction="sum"
        )
        return {
            "loss_fcos_cls": loss_cls / normalizer,
            "loss_fcos_loc": loss_box_reg / normalizer,
            "loss_fcos_ctr": ctrness_loss / normalizer,
        }

    def compute_ctrness_targets(self, anchors: List[Boxes], gt_boxes: List[torch.Tensor]):
        anchors = Boxes.cat(anchors).tensor  # Rx4
        reg_targets = [self.box2box_transform.get_deltas(anchors, m) for m in gt_boxes]
        reg_targets = torch.stack(reg_targets, dim=0)  # NxRx4
        if len(reg_targets) == 0:
            return reg_targets.new_zeros(len(reg_targets))
        left_right = reg_targets[:, :, [0, 2]]
        top_bottom = reg_targets[:, :, [1, 3]]
        ctrness = (left_right.min(dim=-1)[0] / left_right.max(dim=-1)[0]) * (
            top_bottom.min(dim=-1)[0] / top_bottom.max(dim=-1)[0]
        )
        return torch.sqrt(ctrness)

    def forward_inference(
        self,
        images: ImageList,
        features: List[torch.Tensor],
        predictions: List[List[torch.Tensor]],
    ):
        pred_logits, pred_anchor_deltas, pred_centerness = self._transpose_dense_predictions(
            predictions, [self.num_classes, 4, 1]
        )
        anchors = self.anchor_generator(features)

        results: List[Instances] = []
        for img_idx, image_size in enumerate(images.image_sizes):
            scores_per_image = [
                # Multiply and sqrt centerness & classification scores
                # (See eqn. 4 in https://arxiv.org/abs/2006.09214)
                torch.sqrt(x[img_idx].sigmoid_() * y[img_idx].sigmoid_())
                for x, y in zip(pred_logits, pred_centerness)
            ]
            deltas_per_image = [x[img_idx] for x in pred_anchor_deltas]
            results_per_image = self.inference_single_image(
                anchors, scores_per_image, deltas_per_image, image_size
            )
            results.append(results_per_image)
        return results

    def inference_single_image(
        self,
        anchors: List[Boxes],
        box_cls: List[torch.Tensor],
        box_delta: List[torch.Tensor],
        image_size: Tuple[int, int],
    ):
        """
        Identical to :meth:`RetinaNet.inference_single_image.
        """
        pred = self._decode_multi_level_predictions(
            anchors,
            box_cls,
            box_delta,
            self.test_score_thresh,
            self.test_topk_candidates,
            image_size,
        )
        keep = batched_nms(
            pred.pred_boxes.tensor, pred.scores, pred.pred_classes, self.test_nms_thresh
        )
        return pred[keep[: self.max_detections_per_image]]

class FCOSHead(RetinaNetHead):
    """
    The head used in :paper:`fcos`. It adds an additional centerness
    prediction branch on top of :class:`RetinaNetHead`.
    """

    def __init__(self, *, input_shape: List[ShapeSpec], conv_dims: List[int], **kwargs):
        super().__init__(input_shape=input_shape, conv_dims=conv_dims, num_anchors=1, **kwargs)
        # Unlike original FCOS, we do not add an additional learnable scale layer
        # because it's found to have no benefits after normalizing regression targets by stride.
        self._num_features = len(input_shape)
        self.ctrness = nn.Conv2d(conv_dims[-1], 1, kernel_size=3, stride=1, padding=1)
        torch.nn.init.normal_(self.ctrness.weight, std=0.01)
        torch.nn.init.constant_(self.ctrness.bias, 0)

    def forward(self, features):
        assert len(features) == self._num_features
        logits = []
        bbox_reg = []
        ctrness = []
        for feature in features:
            logits.append(self.cls_score(self.cls_subnet(feature)))
            bbox_feature = self.bbox_subnet(feature)
            bbox_reg.append(self.bbox_pred(bbox_feature))
            ctrness.append(self.ctrness(bbox_feature))
        return logits, bbox_reg, ctrness


"""Sample prediction script for ONNX Runtime"""
import argparse
import onnxruntime
import numpy as np
from PIL import Image

MODEL_FILENAME = 'model.onnx'
LABELS_FILENAME = 'labels.txt'

od_model = None

class ObjectDetection:
    INPUT_TENSOR_NAME = 'data'
    OUTPUT_TENSOR_NAMES = ['classLabel', 'loss']
    def __init__(self, model_filename):
        self.session = onnxruntime.InferenceSession(model_filename)
        self.input_shape = self.session.get_inputs()[0].shape[2:]
        self.is_fp16 = self.session.get_inputs()[0].type == 'tensor(float16)'

    def crop_center(self, pil_img, crop_width, crop_height):
        img_width, img_height = pil_img.size
        return pil_img.crop(((img_width - crop_width) // 2,
                            (img_height - crop_height) // 2,
                            (img_width + crop_width) // 2,
                            (img_height + crop_height) // 2))
    def crop_max_square(self,pil_img):
        return self.crop_center(pil_img, min(pil_img.size), min(pil_img.size))

    def predict_image(self, image):
        image = image.convert('RGB') if image.mode != 'RGB' else image
        image = self.crop_max_square(image)
        data = np.asarray(image)
        image = Image.fromarray(np.roll(data,1,axis=-1))
        image = image.resize(self.input_shape)

        inputs = np.array(image, dtype=np.float32)[np.newaxis, :, :, :]
        inputs = inputs.transpose((0, 3, 1, 2))

        if self.is_fp16:
            inputs = inputs.astype(np.float16)

        outputs = self.session.run(self.OUTPUT_TENSOR_NAMES, {self.INPUT_TENSOR_NAME: inputs})
        return (outputs)

def initialize(model_filename):
    global od_model
    od_model = ObjectDetection(model_filename)

def predict(image):
    return od_model.predict_image(image)


import arcade.key
import random
import math
from random import randint
from models import Player
SCREEN_WIDTH = 480
SCREEN_HEIGHT = 640
MOVEMENT_SPEED = 2
LASER_SPEED = 10
SPRITE_SCALING = 0.7

class Laser_E(arcade.Sprite):
    def setup(self, x, y):
        self.center_x = x
        self.top = y
    def update(self):
        self.center_y -= LASER_SPEED-8

class Enemy1(arcade.Sprite):
    def setup(self, x, y,  laser_e_list):
        self.laser_e_list = laser_e_list
        self.center_x = x
        self.center_y = y
        choice_y = [475,495,525,555]
        self.pos_y = choice_y.pop(random.randrange(len(choice_y)))
        self.frame_count = 0

    def update(self,delta):
        self.frame_count += 1
        if self.center_y > self.pos_y:
            self.center_y -= 3

class Enemy2(arcade.Sprite):
    def setup(self, x, y,  laser_e_list):
        self.laser_e_list = laser_e_list
        self.center_x = x
        self.center_y = y
        choice_y = [495,525]
        self.pos_y = choice_y.pop(random.randrange(len(choice_y)))
        self.frame_count = 0

    def update(self,delta):
        self.frame_count += 1
        if self.center_y > self.pos_y:
            self.center_y -= 3

class Boss1(arcade.Sprite):
    def setup(self, x, y,  boss_beam, beam_state):
        self.boss_beam = boss_beam
        self.beam_state = beam_state
        self.center_x = x
        self.center_y = y
        self.pos_y = 550
        self.frame_count = 0

    def update(self,delta):
        self.frame_count += 1
        self.center_x += self.change_x
        if self.center_y > self.pos_y:
            self.center_y -= 3


import datetime
import time

from utils.dictionary import EntityDictionary
from utils.mention import extract_mention_and_plain_text_from_annotated_doc

def extract_mention_and_out_links_from_corpus(corpus_path):
    """
        只得到 mention_anchors 和 out_links，不需要同步生成 train_text
        由于中文 train_text 的生成需要分词，分词很耗时，可以先用这个函数生成一份 mention_anchors 和 out_links

    :param corpus_path:
    :return:
    """
    mention_anchors = dict()
    out_links = dict()
    self_links = dict()

    counter, mode_cnt = 0, 1000000
    start_time = int(time.time())
    last_update = start_time
    print("Extracting mention anchors and out links from corpus: \n\t{}".format(corpus_path))
    with open(corpus_path, "r", encoding="utf-8") as rf:
        for line in rf:
            counter += 1
            if counter % mode_cnt == 0:
                curr_update = int(time.time())
                print("\t#{}, batch_time: {}, total_time: {}".format(
                    counter,
                    str(datetime.timedelta(seconds=curr_update-last_update)),
                    str(datetime.timedelta(seconds=curr_update-start_time))
                ))
                last_update = curr_update
            try:
                instance_id, document = line.strip().split("\t\t")
                mention_anchor_list, _ = extract_mention_and_plain_text_from_annotated_doc(document)
                if out_links.get(instance_id) is None:
                    out_links[instance_id] = set()
                for mention, anchor, offset in mention_anchor_list:
                    mention = mention.lower()
                    if mention_anchors.get(mention) is None:
                        mention_anchors[mention] = dict()
                    if mention_anchors[mention].get(anchor) is None:
                        mention_anchors[mention][anchor] = 0
                    mention_anchors[mention][anchor] += 1
                    out_links[instance_id].add(anchor)

                    # 2020.10.28
                    if (instance_id == anchor):
                        self_links[mention] = self_links.get(mention, 0)+1

            except Exception as e:
                print(counter, e)
    ol = dict()
    for i in out_links:
        if len(out_links[i]) > 0:
            ol[i] = list(out_links[i])
    print("Extracted, total mentions: #{}, total time: {}".format(
        len(mention_anchors), str(datetime.timedelta(seconds=int(time.time())-start_time))))
    return mention_anchors, ol, self_links

def merge_mention_anchors(mention_anchors_list):
    """
        把多源的 mention_anchors 合并起来，例如合并分别从 abstract, article, infobox 中抽取的 mention_anchors

    :param mention_anchors_list:
    :return:
    """
    print("\nMerging mention anchors from {} sources".format(len(mention_anchors_list)))
    start_at = int(time.time())
    ma = dict()
    for mention_anchors in mention_anchors_list:
        for mention in mention_anchors:
            if len(mention) <= 1: continue
            if ma.get(mention) is None:
                ma[mention] = dict()
            for anchor in mention_anchors[mention]:
                if ma[mention].get(anchor) is None:
                    ma[mention][anchor] = 0
                ma[mention][anchor] += mention_anchors[mention][anchor]
    print("Merged, mentions: #{}, time: {}".format(
        len(ma), str(datetime.timedelta(seconds=int(time.time())-start_at))))
    return ma

def merge_out_links(out_links_list):
    """
        把多源的 out_links 合并起来，例如合并分别从 abstract, article, infobox 中抽取的 out_links

    :param out_links_list:
    :return:
    """
    print("\nMerging out links from {} sources".format(len(out_links_list)))
    start_at = int(time.time())
    ol = dict()
    for out_links in out_links_list:
        for inst in out_links:
            if len(out_links[inst]) > 0:
                if ol.get(inst) is None:
                    ol[inst] = set()
                for out_inst in out_links[inst]:
                    ol[inst].add(out_inst)
    for i in ol:
        if len(ol[i]) > 0:
            ol[i] = list(ol[i])
    print("Merged, out_links: #{}, time: {}".format(
        len(ol), str(datetime.timedelta(seconds=int(time.time()) - start_at))))
    return ol

def merge_self_links(self_links_list):
    """
        把多源的 self_links 合并起来，例如合并分别从 abstract, article, infobox 中抽取的 self_links

    :param self_links_list:
    :return:
    """
    print("\nMerging self links from {} sources".format(len(self_links_list)))
    start_at = int(time.time())
    sl = dict()
    for self_links in self_links_list:
        for inst in self_links:
            sl[inst] = sl.get(inst, 0) + self_links[inst]

    print("Merged, self_links: #{}, time: {}".format(
        len(sl), str(datetime.timedelta(seconds=int(time.time()) - start_at))))
    return sl

def expand_mention_anchors(source, mention_anchors):
    """
    从 mention_anchor.json 扩充词典
        a. 将满足以下条件的实体加入到全文统计的实体中，出现次数记为 1
            - 其 title 与 mention-anchor 字典中的某一 mention 相同
            - 该实体从未在语料中以 title 作为 mention 出现过
        b. 对于 title 没有作为 mention 出现过的实体
            - 以 title 作为 mention 构造 title-entity 字典

    :param source: string
    :param mention_anchors: dict
    :return: (dict, dict)
    """
    entity_dict = EntityDictionary.get_instance(source)

    title_entities = dict()

    print("\nExpanding mention anchors from entity dictionary...")
    start_at = int(time.time())
    for instance_id in entity_dict.entity_dict:
        mention = entity_dict.get_entity_by_id(instance_id).get_mention()
        if mention_anchors.get(mention) is not None:
            if mention_anchors[mention].get(instance_id) is None:
                mention_anchors[mention][instance_id] = 1
        else:
            title_entities[mention] = instance_id
    print("Expanded, title entities: #{}, mentions: #{}, time: {}".format(
        len(title_entities), len(mention_anchors), str(datetime.timedelta(seconds=int(time.time())-start_at))))
    return title_entities

def filter_mention_anchors(mention_anchors, link_m, freq_m, self_links, link_prob_th) -> dict:
    """
    1. filter out len(m) <= 1
    2. expand mention_anchors from entity dictionary
    3. filter out link(m) < 2
    4. filter out link_prob(m) < 0.0001
    :param mention_anchors:
    :return:
    """
    ma = dict()
    for m in mention_anchors:
        if len(m) > 1:
            ma[m] = mention_anchors[m]
    nma = dict()
    for m in ma:
        if m == '__all__':
            continue
        if link_m.get(m) is None or freq_m.get(m) is None or (link_m[m] - self_links.get(m, 0) < 2):
            continue
        if (float(link_m[m])/float(freq_m[m])) < link_prob_th:
            continue
        nma[m] = ma[m]
    return nma


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import asyncio
import datetime
import logging
import time
import traceback
import typing

import asyncpg
from discord.ext import commands

import base
from base.utils.provably_fair import update_seed
from . import config, database

class Client(commands.Bot):
    """
    Based off of commands.Bot, with a built in connection pool, amongst other things.
    """

    def __init__(self, configuration=config.Config) -> None:
        self.server_seed_hash = str
        self.server_seed = str
        self.nonce = int
        self.open_games = dict()
        self.game_numbers = set()
        self.raffles = dict()

        self.started_at = float("nan")
        self.logger = logging.getLogger(__name__)
        self.config = configuration
        self.database: typing.Optional[asyncpg.pool.Pool] = None
        self.sql_cache = database.SQLCache()
        self.command_invoke_count = 0
        super().__init__(command_prefix=self.config.bot.command_prefix)

    @property
    def uptime(self) -> datetime:
        return datetime.timedelta(seconds=time.perf_counter() - self.started_at)

    async def start(self) -> None:
        self.database = await database.create_connection_pool(self.sql_cache, self.config.postgres)
        self._load_all_extensions()
        self.logger.info("Proceeding with startup of bot")

        self.started_at = time.perf_counter()
        self.loop.create_task(update_seed(self))
        await self._start()

    async def _start(self) -> None:
        await super().start(self.config.bot.token)

    async def close(self) -> None:
        self.logger.info("Closing asyncpg connection")
        try:
            if self.database is not None:
                await self.database.close()
        finally:
            await super().close()

    def _load_all_extensions(self):
        self.logger.info("Loading extensions...")
        extensions = set(base.extensions) - set(self.config.bot.blacklist_extensions)

        for ext in extensions:
            try:
                self.load_extension("base.exts." + ext, False)
            except Exception as ex:
                self.logger.warning("Failed to load %s because of exception", ext, exc_info=ex)
            else:
                self.logger.info("Loaded extension %s successfully", ext)

    def load_extension(self, name, mute=True):
        # Annoyingly, Dpy calls load_extension meaning we cant attach a logger to it directly or we get
        # two log entries...
        if not mute:
            self.logger.info("Loading %s", name)
        super().load_extension(name)

    def unload_extension(self, name):
        self.logger.info("Unloading %s", name)
        super().unload_extension(name)

    def reload_extension(self, name):
        self.logger.info("Reloading %s", name)
        super().reload_extension(name)

    async def on_connect(self):
        await self.get_owner()

    async def on_command_error(self, ctx, ex):
        if isinstance(ex, commands.CommandOnCooldown):
            self.logger.debug("%s is on cool down for %ss", ctx.author, ex.retry_after)
            await ctx.message.add_reaction("\N{SNOWFLAKE}")
            await asyncio.sleep(ex.retry_after)
            await ctx.message.remove_reaction("\N{SNOWFLAKE}", self.user)
        else:
            self.logger.error("".join(traceback.format_exception(type(ex), ex, ex.__traceback__, 4)))

    async def on_command(self, ctx):
        self.command_invoke_count += 1

    async def get_owner(self):
        if self.owner_id is None:
            info = await self.application_info()
            self.owner_id = info.owner.id
        return self.owner_id


import numpy as np
from utils import load_obj,split,textprocess
import os
import onnxruntime
from metrics import get_entities
import json

def get_attn_pad_mask(seq_q, seq_k):
    # print(seq_q)
    batch_size = 1
    len_q = len(seq_q[0])
    len_k = len(seq_k[0])
    # eq(zero) is PAD token
    pad_attn_mask = np.array(seq_k)==0  # (batch_size, 1, len_k/len_q) one is masking
    # pad_mask = np.repeat(pad_attn_mask,len_q,axis=0).reshape(1,len_q,len_k)
    return pad_attn_mask
    # return pad_attn_mask.repeat(batch_size, len_q, len_k)  # (batch_size, len_q, len_k)

def softmax(x):
    """
    Compute the softmax function for each row of the input x.

    Arguments:
    x -- A N dimensional vector or M x N dimensional numpy matrix.

    Return:
    x -- You are allowed to modify x in-place
    """
    orig_shape = x.shape

    if len(x.shape) > 1:
        # Matrix
        exp_minmax = lambda x: np.exp(x - np.max(x))
        denom = lambda x: 1.0 / np.sum(x)
        x = np.apply_along_axis(exp_minmax,1,x)
        denominator = np.apply_along_axis(denom,1,x)

        if len(denominator.shape) == 1:
            denominator = denominator.reshape((denominator.shape[0],1))

        x = x * denominator
    else:
        # Vector
        x_max = np.max(x)
        x = x - x_max
        numerator = np.exp(x)
        denominator =  1.0 / np.sum(numerator)
        x = numerator.dot(denominator)

    assert x.shape == orig_shape
    return x

def softmax_mask(logits_tgt, cls_idx, idx_mask):
    logits_tgt = logits_tgt[0,:,:]
    mask = idx_mask[str(cls_idx)]
    length, tgt_num = logits_tgt.shape[0], logits_tgt.shape[1]
    scores_exp = softmax(logits_tgt)
     # this step masks
    scores_exp = scores_exp * mask
    return scores_exp

class DataLoader_test(object):
    def __init__(self, save_dir):
        self.save_dir = save_dir
        self.word2idx = load_obj(self.save_dir + "dict.json")
        self.config = load_obj(self.save_dir + "Config.json")
        self.max_len = self.config["max_len"]

        self.WORD = {int(k):v for k,v in self.config["WORD"].items()}
        self.BOS = self.config["BOS"]
        self.UNK = self.config["UNK"]
        self.PAD = self.config["PAD"]

        assert self.BOS == self.word2idx[self.WORD[self.BOS]]
        assert self.UNK == self.word2idx[self.WORD[self.UNK]]
        assert self.PAD == self.word2idx[self.WORD[self.PAD]]

    def load_sentences(self, sent):
        """Loads sentences and tags from their corresponding files.
            Maps tokens and tags to their indices and stores them in the provided dict d.
        """
        sentence = []

        tokens = split(textprocess(sent))
        sentence.append(self.convert_tokens_to_ids(tokens))
        return tokens, sentence

    def convert_tokens_to_ids(self, tokens):
        sentence = []
        sentence.append(self.BOS)
        for tok in tokens:
            if tok in self.word2idx:
                sentence.append(self.word2idx[tok])
            else:
                sentence.append(self.UNK)
        pad = [self.PAD]*(self.max_len+1 - len(sentence))

        assert len(sentence + pad) == self.max_len+1
        return sentence + pad

class NLU_module():
    def __init__(self, save_dir = "./model_onnx/",model_nm = "transformer_mix.onnx"):
        self.save_dir = save_dir
        self.model_nm = model_nm
        self.Init_model()

    def Init_model(self):
        #init dataloader
        self.data_loader = DataLoader_test(self.save_dir)
        # init model
        self.ort_session = onnxruntime.InferenceSession(self.save_dir+ self.model_nm)
        # init dict
        self.idx2lbl = load_obj(self.save_dir + "idx2lbl.json")
        self.idx2cls  = load_obj(self.save_dir + "idx2cls.json")
        # get valid slot for a specific intent
        self.idx_mask = load_obj(self.save_dir + "idx_mask_onnx.json")

    def Inference(self, input_sentence):
        # read test_sentence
        # input_sentence = '导航到世纪大道一百一十八号'
        tokens, test_data = self.data_loader.load_sentences(input_sentence)

        # run inference
        pred_cls ,pred_lbls = self.test(test_data)

        # merge_slot
        slot = self.merged_slot(tokens, pred_lbls)

        ans = {}
        ans["Input_sentence"] = input_sentence.encode('utf-8').decode('utf-8')
        ans["Raw Labels"] = ' '.join(pred_lbls)
        ans["Intent"] = ''.join(pred_cls)
        ans["Megred Mentions"] = slot

        # return json.dumps(ans, ensure_ascii=False)
        return ans

    def test(self, enc):
        """Evaluate the model on `steps` batches."""
        enc_self_attn_mask = get_attn_pad_mask(enc, enc)
        x = (enc,enc_self_attn_mask)

        # compute ONNX Runtime output prediction
        ort_inputs = {self.ort_session.get_inputs()[i].name: x[i] for i in range(len(x))}
        logits_tgt, logits_clsf = self.ort_session.run(None, ort_inputs)

        pad_num = np.count_nonzero(enc)

        pred_cls = np.argmax(logits_clsf)
        score = logits_clsf[0,pred_cls]

        masked_logits_tgt= softmax_mask(logits_tgt, pred_cls, self.idx_mask)
        tgt_idx = np.argmax(masked_logits_tgt, axis = 1)
        score_tgt = 0

        pred_tags = tgt_idx[:pad_num]

        pred_lbls = []
        for idx in pred_tags:
            pred_lbls.append(self.idx2lbl[str(idx)])
        pred_cls = self.idx2cls[str(pred_cls)]

        return pred_cls ,pred_lbls

    def merged_slot(self,tokens, pred_lbls):
        chunks = get_entities(pred_lbls)
        slot_result = {}
        for chunk in chunks:
            tag, start, end = chunk[0], chunk[1], chunk[2]
            tok = ''.join(tokens[chunk[1]:chunk[2]+1])
            # string = '<{0}>: {1}'.format(tag, tok)
            while tag in slot_result:
                tag += '#'
            slot_result[tag]=tok
        return slot_result

if __name__ == '__main__':
    Module = NLU_module()
    # read test_sentence
    input_sentence = '找个北京的餐馆'
    results = Module.Inference(input_sentence)
    print(results)


from __future__ import unicode_literals
from input_reader import InputReader, ReaderError
from pytest import raises

# Subclass InputReader to implement the post_process method
class MyInputReader(InputReader):
    def __init__(self):
        super(MyInputReader, self).__init__()
    def post_process(self, namespace):
        """ Process the regular expression in a way that is easier to use"""

        allowedcolors = ('red', 'green', 'blue', 'yellow', 'violet')
        numbers = []
        colors = []
        for r in namespace.sample:
            # Check validity of numbers
            num = float(r.group(1))
            if not (1000 > num > -1000):
                raise ReaderError ('sample: given number range must be '
                                   '-1000 < num < 1000, given '+str(num))
            numbers.append(num)

            # Check validity of colors
            col = r.group(2)
            if col not in allowedcolors:
                c = ', '.join([repr(x) for x in allowedcolors[:-2]])
                raise ReaderError ('sample: allowed colors are '+c+
                              ' '+repr(allowedcolors[-1])+', given '+repr(col))
            else:
                colors.append(col)

        # Add the results to the namespace
        namespace.add('numbers', tuple(numbers))
        namespace.add('colors', tuple(colors))

def test_custom_reader_works():
    # Use the custom input reader
    reader = MyInputReader()
    reader.add_regex_line('sample', r'(-?\d+\.?\d*) (\w+)', repeat=True)

    inp = reader.read_input(['40.432 red',
                             '-593 blue'])

    assert inp.sample[0].group(0) == '40.432 red'
    assert inp.sample[1].group(0) == '-593 blue'
    assert inp.numbers == (40.432, -593.0)
    assert inp.colors == ('red', 'blue')

def test_custom_reader_error():
    # Errors
    reader = MyInputReader()
    reader.add_regex_line('sample', r'(-?\d+\.?\d*) (\w+)', repeat=True)

    with raises(ReaderError) as e:
        inp = reader.read_input(['40.432 black'])
    assert 'allowed colors are' in str(e.value)

    with raises(ReaderError) as e:
        inp = reader.read_input(['-2000 blue'])
    assert ' given number range must be -1000 < num < 1000' in str(e.value)


from  django import forms
from .models import ItemInformation,ItemCategory,ItemPackSize,ItemBrand,UnitofMeasure
from parsley.decorators import parsleyfy
from django.forms  import ModelForm

@parsleyfy
class ItemInformationForm(ModelForm):
    class Meta:
        model = ItemInformation

        labels = {
            'ItemName': 'Item Name : ',
            'ItemDescription': 'Item Description : ',
            'ItemCode': 'Item Code : ',
            'IsActive':'Status : ',
            'Remarks': 'Remarks : ',
            'ItemPackSize': 'Item Size : ',
            'ItemBrand': 'Brand Name : ',
            'ItemCategory': 'Item Category : ',
            'Uom': 'UOM : '
        }
        fields = ['ItemName','ItemCode','ItemPackSize','ItemBrand','ItemCategory','Uom','ItemDescription','Remarks','IsActive','CreatedBy','CreatedDate','UpdatedBy','UpdatedDate','DeletedBy','DeletedDate' ]
        exclude = ('CreatedBy','CreatedDate','UpdatedBy','UpdatedDate','DeletedBy','DeletedDate')
        widgets = {
            'ItemName': forms.TextInput(),
            'ItemDescription': forms.Textarea(attrs={'rows':2}),
            'ItemCode': forms.TextInput(),
            'IsActive': forms.CheckboxInput(),
            'Remarks': forms.Textarea(attrs={'rows':2}),
            'ItemPackSize': forms.Select(),
            'ItemBrand': forms.Select(),
            'ItemCategory': forms.Select(),
            'Uom': forms.Select()
        }
        parsley_extras = {
            "ItemName": {
                "required-message": "Item Name Required",
            },
            "ItemDescription":{
                "required-message":"Item Description Required"
            },
            "ItemCode":{
                "required-message":"Item Code Required"
            }
        }
    def __init__(self, *args, **kwargs):
        super(ItemInformationForm, self).__init__(*args, **kwargs)
        self.fields['ItemName'].widget.attrs['class'] = 'form-control description'
        self.fields['ItemDescription'].widget.attrs['class'] = 'form-control ItemDescription'
        self.fields['ItemDescription'].widget.attrs['readonly'] = True # text input
        #myform.fields['status'].widget.attrs['disabled'] = True # radio / checkbox
        self.fields['ItemCode'].widget.attrs['class'] = 'form-control description'
        self.fields['IsActive'].widget.attrs['class'] = 'iCheck'
        self.fields['Remarks'].widget.attrs['class'] = 'form-control'
        self.fields['ItemPackSize'].widget.attrs['class'] = 'form-control select2 description'
        self.fields['ItemPackSize'].queryset = ItemPackSize.objects.filter(IsActive = True,DeletedBy=None).distinct()
        self.fields['ItemBrand'].widget.attrs['class'] = 'form-control select2'
        self.fields['ItemBrand'].queryset = ItemBrand.objects.filter(IsActive = True,DeletedBy=None)
        self.fields['ItemCategory'].widget.attrs['class'] = 'form-control select2'
        self.fields['ItemCategory'].queryset = ItemCategory.objects.filter(IsActive = True,DeletedBy=None)
        self.fields['Uom'].widget.attrs['class'] = 'form-control select2'
        self.fields['Uom'].queryset = UnitofMeasure.objects.filter(IsActive = True,DeletedBy=None)
        # self.fields['Uom'].queryset = UnitofMeasure.objects.filter(IsActive = True,DeletedBy=None)


import os
import pickle
from shutil import copyfile
from datetime import datetime
from tqdm import tqdm
import torch
from torch.utils.data import DataLoader
from torch.cuda import get_device_name, current_device
import torch.nn.utils.rnn as rnn
# from tensorboardX import SummaryWriter
from vqa_dataset import VQADataSet
from model import CoattentionNet
from image_encoder import load_resnet

class ExperimentRunner():
    ''' Base class for runnung the experiment.
    '''

    def __init__(self, train_images_dir, train_questions_dir, train_annotations_dir,
                val_images_dir, val_questions_dir, val_annotations_dir , batch_size,
                num_of_epochs, num_of_workers, collate, train_results_dir, val_results_dir,
                saving_frequency, learning_rate, save_results=False):
        ''' Initializes the experiment runner class.
        Parameters:
            train_images_dir: string; Path to train images.
            train_questions_dir: string; Path to train questions.
            train_annotations_dir: string; Path to train annotations.
            val_images_dir: string; Path to validation images.
            val_questions_dir: string; Path to validation questions.
            val_annotations_dir: string; Path to validation annotations.
            batch_size: int; Batch size.
            num_of_epochs: int; Number of epochs.
            num_of_workers: int; Number of workers.
            collate: boolean; Flag to indicate that the results have already been preprocessed and saved. We will not need to do it again. Also images have been encoded.
            train_results_dir: string; Path to the saved train results.
            val_results_dir: string; Path to the saved validation results.
            saving_frequency: int; Frequency of saving in checkpoints.
            learning_rate: float; The learning rate of the algorithm.
            save_results: boolean; Flag for saving results such as vocabulary, top 1000 answers etc.
        '''

        self.train_dataset = VQADataSet("train", train_questions_dir, train_annotations_dir, train_images_dir,
                                  "COCO_train2014_", collate, save_results, train_results_dir)
        self.val_dataset = VQADataSet("val", val_questions_dir, val_annotations_dir, val_images_dir,
                                "COCO_val2014_", collate, save_results, val_results_dir)

        self.train_dataloader = DataLoader(self.train_dataset, batch_size=batch_size, shuffle=True,
                                          num_workers=num_of_workers, collate_fn=self.custom_collate_func)
        self.val_dataloader = DataLoader(self.val_dataset, batch_size=batch_size, shuffle=False,
                                        num_workers=num_of_workers, collate_fn=self.custom_collate_func)

        self.num_of_epochs = num_of_epochs
        self.num_of_workers = num_of_workers
        self.logging_frquency = 10  # every 10 training steps
        self.verbose = 50 # print every 50 steps
        self.batch_size = batch_size
        self.saving_frequency = saving_frequency  # epoch-wise
        self.learning_rate = learning_rate

        # Loading the model
        with open(os.path.join(train_results_dir, "question_vocabulary.pkl"), 'rb') as f:
            question_vocabulary = pickle.load(f)

        self.model = CoattentionNet(len(question_vocabulary), 1000).float()
        self.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
        if self.DEVICE == "cuda":
            print(f"Loading model with: {self.DEVICE} | {get_device_name(current_device())}...")
            self.model = self.model.cuda()
        else:
            print(f"Loading model with: {self.DEVICE}...")

        # Setting the optimizer
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.learning_rate, weight_decay=1e-8)

        # Setting the loss function
        self.criterion = torch.nn.CrossEntropyLoss()

        # Weights initialization
        self.initialize_weights()

        # Logger
        # self.logger = SummaryWriter()  # rewrite this.

        # Image Encoder
        self.image_encoder = load_resnet(self.DEVICE)

        # Checkpoints
        self.checkpoints_dir = os.path.join("Data", "saved")
        if not os.path.isdir(self.checkpoints_dir):
            os.makedirs(self.checkpoints_dir)
        self.experiment_id = len(os.listdir(self.checkpoints_dir))

        if not os.path.isdir(os.path.join(self.checkpoints_dir, str(self.experiment_id))):
            print(f"Creatting {os.path.join(self.checkpoints_dir, str(self.experiment_id))}...")
            os.makedirs(os.path.join(self.checkpoints_dir, str(self.experiment_id)))

    def optimize(self, predicted_answers, true_answers):
        ''' Optimization step for the model.
        Parameters:
            predicted_answers: pytorch tensor object; Tensor of the model predictions.
            true_answers: pytorch tensor object; Ids of ground truth answers.
        Returns:
            loss: float; Loss for the model in current step.
        '''

        self.optimizer.zero_grad()
        loss = self.criterion(predicted_answers, true_answers)
        loss.backward()
        self.optimizer.step()

        return loss

    def validate(self):
        ''' Method for validating the model after training.
        Returns:
            accuracy: float; Accuracy of the model after validation.
        '''

        print("Validating the model...")
        accuracy = 0.0
        progress_bar = tqdm(total=len(self.val_dataloader))
        self.model.eval()
        with torch.no_grad():
            for batch_idx, (image_tensors, question_tensors, answer_tensors) in enumerate(self.val_dataloader):
                progress_bar.update(1)
                question_tensors = rnn.pack_sequence(question_tensors)
                image_tensors = image_tensors.to(self.DEVICE)
                image_tensors = self.image_encoder(image_tensors)
                image_tensors = image_tensors.view(image_tensors.size(0), image_tensors.size(1), -1)

                question_tensors = question_tensors.to(self.DEVICE)
                answer_tensors = answer_tensors.to(self.DEVICE)

                answer_tensors = torch.squeeze(answer_tensors)

                # predictions
                predicted_answer = self.model(image_tensors, question_tensors)

                for i in range(answer_tensors.shape[0]):
                    if torch.argmax(predicted_answer[i]).item() == answer_tensors[i]:
                        accuracy += 1.0

                # if (batch_idx + 1) % self.verbose == 0:
                #     print(f"validation accuracy: {round(accuracy / ((batch_idx + 1) * self.batch_size), 2)}")

            accuracy = round(accuracy / len(self.val_dataset), 2)

            return accuracy

    def train(self):
        ''' Method for training the model.
        '''

        print("\n\n--------------------------------------------------\nTraining the model...\n--------------------------------------------------\n\n")
        train_iteration = 0
        val_iteration = 0
        best_accuracy = 0.0
        loss_history = []
        accuracy_history = []

        for epoch in range(self.num_of_epochs):
            progress_bar = tqdm(total=len(self.train_dataloader))
            self.model.train()

            # if (epoch + 1) % 5 == 0:  # HYPER PARAMETER
            #     self.adjust_learning_rate()

            for batch_idx, (image_tensors, question_tensors, answer_tensors) in enumerate(self.train_dataloader):
                progress_bar.update(1)

                current_step = epoch * len(self.train_dataloader) + batch_idx

                question_tensors = rnn.pack_sequence(question_tensors)
                image_tensors = image_tensors.to(self.DEVICE)
                image_tensors = self.image_encoder(image_tensors)
                image_tensors = image_tensors.view(image_tensors.size(0), image_tensors.size(1), -1)

                question_tensors = question_tensors.to(self.DEVICE)
                answer_tensors = answer_tensors.to(self.DEVICE)

                answer_tensors = torch.squeeze(answer_tensors)

                # predictions
                predicted_answer = self.model(image_tensors, question_tensors)

                # Optimize the model according to the predictions
                loss = self.optimize(predicted_answer, answer_tensors)
                loss_history.append(loss)

                if (current_step + 1) % self.logging_frquency == 0:
                    # print(f"Epoch: {epoch}, Batch: {batch_idx}/{len(self.train_dataloader)} has loss: {loss}")

                    # self.logger.add_scalar("train/loss", loss.item(), train_iteration)
                    train_iteration += 1

            print(f"Epoch: {epoch} has loss: {loss}")

            # Validating
            if (epoch + 1) % self.saving_frequency == 0 or epoch == self.num_of_epochs - 1:
                validation_accuracy = self.validate()
                accuracy_history.append(validation_accuracy)
                print(f"Epoch: {epoch} has validation accuracy {validation_accuracy}")

                # self.logger.add_scalar("valid/accuracy", validation_accuracy, val_iteration)
                val_iteration += 1

                # Remember the best validation accuracy and save a checkpoint
                is_best = validation_accuracy > best_accuracy
                best_accuracy = max(best_accuracy, validation_accuracy)
                self.save_checkpoint(epoch, is_best, best_accuracy)

        # Closing tensorboard logger
        # logger_dir = os.path.join(self.checkpoints_dir, str(self.experiment_id), datetime.now().strftime("%d-%m-%y_%H-%M-%S"))
        # if not os.path.isdir(logger_dir):
        #     os.makedirs(logger_dir)
        # self.logger.export_scalars_to_json(logger_dir + 'tensorboard_summary.json')
        # self.logger.close()

    def custom_collate_func(self, seq_list):
        ''' Custome collate function to stack multiple images/questions/answers in a batch.
        Parameters:
            seq_list: *args; Whatever the dataloader returns.
        Returns:
            image_tensors: pytorch tensor object; Stacked image tensor.
            question_tensors: pytorch tensor object; Stacked question tensor.
            answer_tensors: pytorch tensor object; Stacked answer tensor.
        '''

        image_tensors, question_tensors, answer_tensors = zip(*seq_list)  # will recieve a batch of data points.
        lens = [len(question) for question in question_tensors]
        seq_order = sorted(range(len(lens)), key=lens.__getitem__, reverse=True)

        image_tensors = torch.stack([image_tensors[i] for i in seq_order])
        question_tensors = [question_tensors[i] for i in seq_order]
        answer_tensors = torch.stack([answer_tensors[i] for i in seq_order])

        return image_tensors, question_tensors, answer_tensors

    def initialize_weights(self):
        ''' Method to initialize model weights.
        '''

        for layer in self.model.modules():
            if isinstance(layer, (torch.nn.Conv1d, torch.nn.Linear)):
                try:
                    torch.nn.init.xavier_normal_(layer.weight)

                    try:
                        torch.nn.init.constant_(layer.bias.data, 0)
                    except:
                        pass
                except:
                    pass

    def adjust_learning_rate(self):
        ''' Sets the learning rate to the initial learning rate decayed by 10.
        '''

        for param_group in self.optimizer.param_groups:
            param_group['lr'] = param_group['lr'] / 10

    def save_checkpoint(self, epoch, is_best, best_accuracy):
        ''' Method for sacving a checkpoint.
        Parameters:
            epoch: int; Epoch number for the checkpoint.
            is_best: boolean; Whether the checkpoint has the best accuracy or not.
            best_accuracy: float; Best accuracy value.
        '''

        state = {"epoch": epoch + 1,
                "model_state_dict": self.model.state_dict(),
                "optimizer_state_dict": self.optimizer.state_dict(),
                "best_accuracy": best_accuracy
                }

        checkpoint_path = os.path.join(self.checkpoints_dir, str(self.experiment_id), f"{epoch + 1}_checkpoint.pt")
        torch.save(state, checkpoint_path)

        if is_best:
            copyfile(checkpoint_path, os.path.join(self.checkpoints_dir, str(self.experiment_id), "best.pt"))  # create a separate copy of the best checkpoint

if __name__ == "__main__":
    train_image_dir = "Data/train/images/train10K"
    train_qjson = "Data/train/questions/train_quest_10K.json"
    train_ajson = "Data/train/annotations/train_ann_10K.json"
    collate = True
    train_results_dir = "Data/train/cache/"

    val_image_dir = "Data/val/images/val3K"
    val_qjson = "Data/val/questions/val_quest_3K.json"
    val_ajson = "Data/val/annotations/val_ann_3K.json"
    collate = True
    val_results_dir = "Data/val/cache/"

    saving_frequency = 10
    learning_rate = 0.001
    batch_size = 100
    num_epochs = 2
    num_workers = 5

    exp_runner = ExperimentRunner(train_image_dir, train_qjson, train_ajson, val_image_dir, val_qjson, val_ajson,
                                batch_size, num_epochs, num_workers, collate, train_results_dir, val_results_dir, saving_frequency, learning_rate)

    # exp_runner.save_checkpoint(100, True, 92.4)
    exp_runner.train()


__metaclass__ = type

# let's try to keep path imports to a minimum...
from os.path import dirname, split as splitpath

import sys
import zipimport
import inspect
import warnings
import re
try:
    import ast
except ImportError:
    ast = None

from zope.interface import Interface, implements

from exocet._filepath import UnlistableError, FilePath
from exocet._zippath import  ZipArchive

from exocet._reflect import namedAny
from exocet._components import registerAdapter

_nothing = object()

PYTHON_EXTENSIONS = ['.py']
OPTIMIZED_MODE = __doc__ is None
if OPTIMIZED_MODE:
    PYTHON_EXTENSIONS.append('.pyo')
else:
    PYTHON_EXTENSIONS.append('.pyc')
PYTHON_EXTENSIONS.extend([".so", ".pyd"])

def _isPythonIdentifier(string):
    """
    cheezy fake test for proper identifier-ness.

    @param string: a str which might or might not be a valid python identifier.

    @return: True or False
    """
    if string == '':
        return True
    try:
        return re.match('[a-zA-Z_][a-zA-Z0-9_]*$', string) is not None
    except TypeError:
        return False

class NotLoadedError(Exception):
    """
    Attempt to access a value that hasn't been loaded yet.

    @since: 10.2
    """

def _isPackagePath(fpath):
    # Determine if a FilePath-like object is a Python package.  TODO: deal with
    # __init__module.(so|dll|pyd)?
    extless = fpath.splitext()[0]
    basend = splitpath(extless)[1]
    return basend == "__init__"

class _ModuleIteratorHelper:
    """
    This mixin provides common behavior between python module and path entries,
    since the mechanism for searching sys.path and __path__ attributes is
    remarkably similar.
    """

    def iterModules(self):
        """
        Loop over the modules present below this entry or package on PYTHONPATH.

        For modules which are not packages, this will yield nothing.

        For packages and path entries, this will only yield modules one level
        down; i.e. if there is a package a.b.c, iterModules on a will only
        return a.b.  If you want to descend deeply, use walkModules.

        @return: a generator which yields PythonModule instances that describe
        modules which can be, or have been, imported.
        """
        yielded = {}
        if not self.filePath.exists():
            return

        for placeToLook in self._packagePaths():
            try:
                children = placeToLook.children()
            except UnlistableError:
                continue

            children.sort()
            for potentialTopLevel in children:
                ext = potentialTopLevel.splitext()[1]
                potentialBasename = potentialTopLevel.basename()[:-len(ext)]
                if ext in PYTHON_EXTENSIONS:
                    # TODO: this should be a little choosier about which path entry
                    # it selects first, and it should do all the .so checking and
                    # crud
                    if not _isPythonIdentifier(potentialBasename):
                        continue
                    modname = self._subModuleName(potentialBasename)
                    if modname.split(".")[-1] == '__init__':
                        # This marks the directory as a package so it can't be
                        # a module.
                        continue
                    if modname not in yielded:
                        yielded[modname] = True
                        pm = PythonModule(modname, potentialTopLevel, self._getEntry())
                        assert pm != self
                        yield pm
                else:
                    if (ext or not _isPythonIdentifier(potentialBasename)
                        or not potentialTopLevel.isdir()):
                        continue
                    modname = self._subModuleName(potentialTopLevel.basename())
                    for ext in PYTHON_EXTENSIONS:
                        initpy = potentialTopLevel.child("__init__"+ext)
                        if initpy.exists():
                            yielded[modname] = True
                            pm = PythonModule(modname, initpy, self._getEntry())
                            assert pm != self
                            yield pm
                            break

    def walkModules(self, importPackages=False):
        """
        Similar to L{iterModules}, this yields self, and then every module in my
        package or entry, and every submodule in each package or entry.

        In other words, this is deep, and L{iterModules} is shallow.
        """
        yield self
        for package in self.iterModules():
            for module in package.walkModules(importPackages=importPackages):
                yield module

    def _subModuleName(self, mn):
        """
        This is a hook to provide packages with the ability to specify their names
        as a prefix to submodules here.
        """
        return mn

    def _packagePaths(self):
        """
        Implement in subclasses to specify where to look for modules.

        @return: iterable of FilePath-like objects.
        """
        raise NotImplementedError()

    def _getEntry(self):
        """
        Implement in subclasses to specify what path entry submodules will come
        from.

        @return: a PathEntry instance.
        """
        raise NotImplementedError()

    def __getitem__(self, modname):
        """
        Retrieve a module from below this path or package.

        @param modname: a str naming a module to be loaded.  For entries, this
        is a top-level, undotted package name, and for packages it is the name
        of the module without the package prefix.  For example, if you have a
        PythonModule representing the 'twisted' package, you could use::

            twistedPackageObj['python']['modules']

        to retrieve the Twisted-supplied version of this module.

        @raise: KeyError if the module is not found.

        @return: a PythonModule.
        """
        for module in self.iterModules():
            if module.name == self._subModuleName(modname):
                return module
        raise KeyError(modname)

    def __iter__(self):
        """
        Implemented to raise NotImplementedError for clarity, so that attempting to
        loop over this object won't call __getitem__.

        Note: in the future there might be some sensible default for iteration,
        like 'walkEverything', so this is deliberately untested and undefined
        behavior.
        """
        raise NotImplementedError()

if ast is not None:
    _iefBase = ast.NodeVisitor
else:
    _iefBase = object

class _ImportExportFinder(_iefBase):
    """
    Find names of imports and exports (via __all__).

    @ivar imports: A set of (source, name) pairs describing an
                   imported name. if the source is None, the name
                   is an import from the top level.

    @ivar exports: A list of names defined in __all__ in this
                   module, or None if the module has no __all__
                   attribute.

    @ivar definedNames: A set of names created by assignment,
                        class definitions, or function definitions
                        at the top level of this module.
    """

    def __init__(self):
        self.imports = set()
        self.exports = None
        self.definedNames = set()

    def visit_Import(self, node):
        """
        Collect names for all import statements.
        """
        for alias in node.names:
            self.imports.add((None, alias.name))

    def visit_ImportFrom(self, node):
        """
        Collect names and source modules from "import x from y" statements.
        """
        if node.names[0].name == "*":
            raise SyntaxError("Code containing 'import *' cannot be statically analyzed.")
        for name in node.names:
            self.imports.add((node.module, name.name))

    def visit_Module(self, node):
        """
        Look for top-level name bindings and __all__ in a module.
        """
        def collectNames(target, value):
            if isinstance(target, ast.Name):
                collectSingleName(value, target.id)
            if isinstance(target, (ast.List, ast.Tuple)):
                if not (isinstance(value, (ast.List, ast.Tuple))
                        and len(value.elts) == len(target.elts)):
                    value = None
                nameNodes = target.elts
                for (i, n) in enumerate(nameNodes):
                    if value is not None:
                        v = value.elts[i]
                    else:
                        v = None
                    collectNames(n, v)

        def collectSingleName(value, name):
            if name == '__all__':
                checkAll(value)
            else:
                self.definedNames.add(name)

        def checkAll(allval):
            if self.exports is not None:
                    raise SyntaxError("__all__ can only be defined once")
            try:
                self.exports = set(ast.literal_eval(allval))
            except ValueError:
                raise SyntaxError("__all__ must only contain literal Python identifier strings")
            for exp in self.exports:
                if not _isPythonIdentifier(exp):
                    raise SyntaxError("__all__ must only contain literal Python identifier strings")
        for stmt in node.body:
            self.visit(stmt)
            if isinstance(stmt, ast.Assign):
                collectNames(stmt.targets[0], stmt.value)
            elif isinstance(stmt, (ast.FunctionDef, ast.ClassDef)):
                self.definedNames.add(stmt.name)

class PythonAttribute:
    """
    I represent a function, class, or other object that is present.

    @ivar name: the fully-qualified python name of this attribute.

    @ivar onObject: a reference to a PythonModule or other PythonAttribute that
    is this attribute's logical parent.

    @ivar name: the fully qualified python name of the attribute represented by
    this class.
    """
    def __init__(self, name, onObject, loaded, pythonValue):
        """
        Create a PythonAttribute.  This is a private constructor.  Do not construct
        me directly, use PythonModule.iterAttributes.

        @param name: the FQPN
        @param onObject: see ivar
        @param loaded: always True, for now
        @param pythonValue: the value of the attribute we're pointing to.
        """
        self.name = name
        self.onObject = onObject
        self._loaded = loaded
        self.pythonValue = pythonValue

    def __repr__(self):
        return 'PythonAttribute<%r>'%(self.name,)

    def isLoaded(self):
        """
        Return a boolean describing whether the attribute this describes has
        actually been loaded into memory by importing its module.
        """
        return self._loaded

    def load(self, default=_nothing):
        """
        Load the value associated with this attribute.

        @return: an arbitrary Python object, or 'default' if there is an error
        loading it.
        """
        if not self.isLoaded():
            mod = self.onObject.load()
            self.pythonValue = getattr(mod, self.name.split('.')[-1], default)
            self._loaded = True
        return self.pythonValue

    def iterAttributes(self):
        """
        Iterate over the attributes of the value named by this
        object. Only works when the module is loaded.

        @raise NotImplementedError: when the module is not loaded.
        """
        if not self.isLoaded():
            raise NotImplementedError("Static inspection of attributes doesn't"
                                      " go beyond the top level of the module.")
        for name, val in inspect.getmembers(self.load()):
            yield PythonAttribute(self.name + '.' + name, self, True, val)

class PythonModule(_ModuleIteratorHelper):
    """
    Representation of a module which could be imported from sys.path.

    @ivar name: the fully qualified python name of this module.

    @ivar filePath: a FilePath-like object which points to the location of this
    module.

    @ivar pathEntry: a L{PathEntry} instance which this module was located
    from.

    @ivar _finder: an L{_ImportExportFinder}, or None. Used to
    discover names used and defined in a module by inspecting its AST.
    """

    def __init__(self, name, filePath, pathEntry):
        """
        Create a PythonModule.  Do not construct this directly, instead inspect a
        PythonPath or other PythonModule instances.

        @param name: see ivar
        @param filePath: see ivar
        @param pathEntry: see ivar
        """
        assert not name.endswith(".__init__")
        self.name = name
        self.filePath = filePath
        self.parentPath = filePath.parent()
        self.pathEntry = pathEntry

        self._finder = None

    def _getEntry(self):
        return self.pathEntry

    def __repr__(self):
        """
        Return a string representation including the module name.
        """
        return 'PythonModule<%r>' % (self.name,)

    def isLoaded(self):
        """
        Determine if the module is loaded into sys.modules.

        @return: a boolean: true if loaded, false if not.
        """
        return self.pathEntry.pythonPath.moduleDict.get(self.name) is not None

    def _maybeLoadFinder(self):
        """
        Scan a module for imports, exports, and attributes.
        """
        if self._finder is None:
            try:
                tree = ast.parse(self.filePath.getContent())
            except TypeError:
                raise ValueError("Static analysis of module attributes can only be done on Python source.")
            self._finder = _ImportExportFinder()
            self._finder.visit(tree)

    def iterAttributes(self):
        """
        List all the attributes defined in this module.

        On Python 2.6 and later this method can list attributes on a
        module without loading it, via AST inspection. On earlier
        Python versions, they must be loaded. This method will use
        inspect.getmembers on all Python versions if the module is loaded.

        @raise NotImplementedError: if this module is not loaded and
        AST inspection is not possible.

        @return: a generator yielding PythonAttribute instances describing the
        attributes of this module.
        """
        if not self.isLoaded():
            if ast is None:
                raise NotImplementedError("Static analysis of module attributes"
                                          " requires the 'ast' module,"
                                          " found in Python 2.6 or later.")
            self._maybeLoadFinder()
            if self._finder.exports:
                attrs = (set([x[1] for x in self._finder.imports]) | self._finder.definedNames) & self._finder.exports
            else:
                attrs = self._finder.definedNames
            for name in attrs:
                yield PythonAttribute(self.name + '.' + name, self, False, _nothing)
        else:
            for name, val in inspect.getmembers(self.load()):
                yield PythonAttribute(self.name + '.' + name, self, True, val)

    def iterImportNames(self):
        """
        List all the fully-qualified names imported in this module.

        @raise NotImplementedError: if this module is not loaded and
        AST inspection is not possible.

        @return: a generator yielding fully qualified name strings for
                 each name imported into this module.

        @since: 10.2
        """
        if ast is None:
            raise NotImplementedError("Static analysis of module attributes"
                                      " requires the 'ast' module, found in"
                                      " Python 2.6 or later.")

        self._maybeLoadFinder()
        for (origin, name) in self._finder.imports:
            if origin is not None:
                yield origin+ '.' + name
            else:
                yield name

    def iterExportNames(self):
        """
        List all the names exported by this module. If the module
        defines __all__ as a list of literal strings, those names will
        be treated as the module's exports. Otherwise, all names
        defined at the top level of the module will be regarded as
        exports.

        @raise NotImplementedError: if this module is not loaded and
        AST inspection is not possible.

        @return: an iterable of names of attributes of this module.

        @since: 10.2
        """
        if ast is None:
            raise NotImplementedError("Static analysis of module attributes"
                                      " requires the 'ast' module, found in"
                                      " Python 2.6 or later.")
        self._maybeLoadFinder()
        return self._finder.exports or self._finder.definedNames

    def isPackage(self):
        """
        Returns true if this module is also a package, and might yield something
        from iterModules.
        """
        return _isPackagePath(self.filePath)

    def load(self, default=_nothing):
        """
        Load this module.

        @param default: if specified, the value to return in case of an error.

        @return: a genuine python module.

        @raise: any type of exception.  Importing modules is a risky business;
        the erorrs of any code run at module scope may be raised from here, as
        well as ImportError if something bizarre happened to the system path
        between the discovery of this PythonModule object and the attempt to
        import it.  If you specify a default, the error will be swallowed
        entirely, and not logged.

        @rtype: types.ModuleType.
        """
        try:
            return self.pathEntry.pythonPath.moduleLoader(self.name)
        except:                 # this needs more thought...
            if default is not _nothing:
                return default
            raise

    def __eq__(self, other):
        """
        PythonModules with the same name are equal.
        """
        if not isinstance(other, PythonModule):
            return False
        return other.name == self.name

    def __ne__(self, other):
        """
        PythonModules with different names are not equal.
        """
        if not isinstance(other, PythonModule):
            return True
        return other.name != self.name

    def walkModules(self, importPackages=False):
        if importPackages and self.isPackage():
            self.load()
        return super(PythonModule, self).walkModules(importPackages=importPackages)

    def _subModuleName(self, mn):
        """
        submodules of this module are prefixed with our name.
        """
        return self.name + '.' + mn

    def _packagePaths(self):
        """
        Yield a sequence of FilePath-like objects which represent path segments.
        """
        if not self.isPackage():
            return
        if self.isLoaded():
            load = self.load()
            if hasattr(load, '__path__'):
                for fn in load.__path__:
                    if fn == self.parentPath.path:
                        # this should _really_ exist.
                        assert self.parentPath.exists()
                        yield self.parentPath
                    else:
                        smp = self.pathEntry.pythonPath._smartPath(fn)
                        if smp.exists():
                            yield smp
        else:
            yield self.parentPath

class PathEntry(_ModuleIteratorHelper):
    """
    I am a proxy for a single entry on sys.path.

    @ivar filePath: a FilePath-like object pointing at the filesystem location
    or archive file where this path entry is stored.

    @ivar pythonPath: a PythonPath instance.
    """
    def __init__(self, filePath, pythonPath):
        """
        Create a PathEntry.  This is a private constructor.
        """
        self.filePath = filePath
        self.pythonPath = pythonPath

    def _getEntry(self):
        return self

    def __repr__(self):
        return 'PathEntry<%r>' % (self.filePath,)

    def _packagePaths(self):
        yield self.filePath

class IPathImportMapper(Interface):
    """
    This is an internal interface, used to map importers to factories for
    FilePath-like objects.
    """
    def mapPath(self, pathLikeString):
        """
        Return a FilePath-like object.

        @param pathLikeString: a path-like string, like one that might be
        passed to an import hook.

        @return: a L{FilePath}, or something like it (currently only a
        L{ZipPath}, but more might be added later).
        """

class _DefaultMapImpl:
    """ Wrapper for the default importer, i.e. None.  """
    implements(IPathImportMapper)
    def mapPath(self, fsPathString):
        return FilePath(fsPathString)
_theDefaultMapper = _DefaultMapImpl()

class _ZipMapImpl:
    """ IPathImportMapper implementation for zipimport.ZipImporter.  """
    implements(IPathImportMapper)
    def __init__(self, importer):
        self.importer = importer

    def mapPath(self, fsPathString):
        """
        Map the given FS path to a ZipPath, by looking at the ZipImporter's
        "archive" attribute and using it as our ZipArchive root, then walking
        down into the archive from there.

        @return: a L{zippath.ZipPath} or L{zippath.ZipArchive} instance.
        """
        za = ZipArchive(self.importer.archive)
        myPath = FilePath(self.importer.archive)
        itsPath = FilePath(fsPathString)
        if myPath == itsPath:
            return za
        # This is NOT a general-purpose rule for sys.path or __file__:
        # zipimport specifically uses regular OS path syntax in its pathnames,
        # even though zip files specify that slashes are always the separator,
        # regardless of platform.
        segs = itsPath.segmentsFrom(myPath)
        zp = za
        for seg in segs:
            zp = zp.child(seg)
        return zp

registerAdapter(_ZipMapImpl, zipimport.zipimporter, IPathImportMapper)

def _defaultSysPathFactory():
    """
    Provide the default behavior of PythonPath's sys.path factory, which is to
    return the current value of sys.path.

    @return: L{sys.path}
    """
    return sys.path

class PythonPath:
    """
    I represent the very top of the Python object-space, the module list in
    sys.path and the modules list in sys.modules.

    @ivar _sysPath: a sequence of strings like sys.path.  This attribute is
    read-only.

    @ivar moduleDict: a dictionary mapping string module names to module
    objects, like sys.modules.

    @ivar sysPathHooks: a list of PEP-302 path hooks, like sys.path_hooks.

    @ivar moduleLoader: a function that takes a fully-qualified python name and
    returns a module, like twisted.python.reflect.namedAny.
    """

    def __init__(self,
                 sysPath=None,
                 moduleDict=sys.modules,
                 sysPathHooks=sys.path_hooks,
                 importerCache=sys.path_importer_cache,
                 moduleLoader=namedAny,
                 sysPathFactory=None):
        """
        Create a PythonPath.  You almost certainly want to use
        modules.theSystemPath, or its aliased methods, rather than creating a
        new instance yourself, though.

        All parameters are optional, and if unspecified, will use 'system'
        equivalents that makes this PythonPath like the global L{theSystemPath}
        instance.

        @param sysPath: a sys.path-like list to use for this PythonPath, to
        specify where to load modules from.

        @param moduleDict: a sys.modules-like dictionary to use for keeping
        track of what modules this PythonPath has loaded.

        @param sysPathHooks: sys.path_hooks-like list of PEP-302 path hooks to
        be used for this PythonPath, to determie which importers should be
        used.

        @param importerCache: a sys.path_importer_cache-like list of PEP-302
        importers.  This will be used in conjunction with the given
        sysPathHooks.

        @param moduleLoader: a module loader function which takes a string and
        returns a module.  That is to say, it is like L{namedAny} - *not* like
        L{__import__}.

        @param sysPathFactory: a 0-argument callable which returns the current
        value of a sys.path-like list of strings.  Specify either this, or
        sysPath, not both.  This alternative interface is provided because the
        way the Python import mechanism works, you can re-bind the 'sys.path'
        name and that is what is used for current imports, so it must be a
        factory rather than a value to deal with modification by rebinding
        rather than modification by mutation.  Note: it is not recommended to
        rebind sys.path.  Although this mechanism can deal with that, it is a
        subtle point which some tools that it is easy for tools which interact
        with sys.path to miss.
        """
        if sysPath is not None:
            sysPathFactory = lambda : sysPath
        elif sysPathFactory is None:
            sysPathFactory = _defaultSysPathFactory
        self._sysPathFactory = sysPathFactory
        self._sysPath = sysPath
        self.moduleDict = moduleDict
        self.sysPathHooks = sysPathHooks
        self.importerCache = importerCache
        self.moduleLoader = moduleLoader

    def _getSysPath(self):
        """
        Retrieve the current value of the module search path list.
        """
        return self._sysPathFactory()

    sysPath = property(_getSysPath)

    def _findEntryPathString(self, modobj):
        """
        Determine where a given Python module object came from by looking at path
        entries.
        """
        topPackageObj = modobj
        while '.' in topPackageObj.__name__:
            topPackageObj = self.moduleDict['.'.join(
                    topPackageObj.__name__.split('.')[:-1])]
        if _isPackagePath(FilePath(topPackageObj.__file__)):
            # if package 'foo' is on sys.path at /a/b/foo, package 'foo's
            # __file__ will be /a/b/foo/__init__.py, and we are looking for
            # /a/b here, the path-entry; so go up two steps.
            rval = dirname(dirname(topPackageObj.__file__))
        else:
            # the module is completely top-level, not within any packages.  The
            # path entry it's on is just its dirname.
            rval = dirname(topPackageObj.__file__)

        # There are probably some awful tricks that an importer could pull
        # which would break this, so let's just make sure... it's a loaded
        # module after all, which means that its path MUST be in
        # path_importer_cache according to PEP 302 -glyph
        if rval not in self.importerCache:
            warnings.warn(
                "%s (for module %s) not in path importer cache "
                "(PEP 302 violation - check your local configuration)." % (
                    rval, modobj.__name__),
                stacklevel=3)

        return rval

    def _smartPath(self, pathName):
        """
        Given a path entry from sys.path which may refer to an importer,
        return the appropriate FilePath-like instance.

        @param pathName: a str describing the path.

        @return: a FilePath-like object.
        """
        importr = self.importerCache.get(pathName, _nothing)
        if importr is _nothing:
            for hook in self.sysPathHooks:
                try:
                    importr = hook(pathName)
                except ImportError:
                    pass
            if importr is _nothing: # still
                importr = None
        return IPathImportMapper(importr, _theDefaultMapper).mapPath(pathName)

    def iterEntries(self):
        """
        Iterate the entries on my sysPath.

        @return: a generator yielding PathEntry objects
        """
        for pathName in self.sysPath:
            fp = self._smartPath(pathName)
            yield PathEntry(fp, self)

    def __getitem__(self, modname):
        """
        Get a python module by its given fully-qualified name.

        @param modname: The fully-qualified Python module name to load.

        @type modname: C{str}

        @return: an object representing the module identified by C{modname}

        @rtype: L{PythonModule}

        @raise KeyError: if the module name is not a valid module name, or no
            such module can be identified as loadable.
        """
        # See if the module is already somewhere in Python-land.
        moduleObject = self.moduleDict.get(modname)
        if moduleObject is not None:
            # we need 2 paths; one of the path entry and one for the module.
            pe = PathEntry(
                self._smartPath(
                    self._findEntryPathString(moduleObject)),
                self)
            path = inspect.getsourcefile(moduleObject)
            if path is None:
                path = moduleObject.__file__
            mp = self._smartPath(path)
            return PythonModule(modname, mp, pe)

        # Recurse if we're trying to get a submodule.
        if '.' in modname:
            pkg = self
            for name in modname.split('.'):
                pkg = pkg[name]
            return pkg

        # Finally do the slowest possible thing and iterate
        for module in self.iterModules():
            if module.name == modname:
                return module
        raise KeyError(modname)

    def __repr__(self):
        """
        Display my sysPath and moduleDict in a string representation.
        """
        return "PythonPath(%r,%r)" % (self.sysPath, self.moduleDict)

    def iterModules(self):
        """
        Yield all top-level modules on my sysPath.
        """
        for entry in self.iterEntries():
            for module in entry.iterModules():
                yield module

    def walkModules(self, importPackages=False):
        """
        Similar to L{iterModules}, this yields every module on the path, then every
        submodule in each package or entry.
        """
        for package in self.iterModules():
            for module in package.walkModules(importPackages=False):
                yield module

theSystemPath = PythonPath()

def walkModules(importPackages=False):
    """
    Deeply iterate all modules on the global python path.

    @param importPackages: Import packages as they are seen.
    """
    return theSystemPath.walkModules(importPackages=importPackages)

def iterModules():
    """
    Iterate all modules and top-level packages on the global Python path, but
    do not descend into packages.

    @param importPackages: Import packages as they are seen.
    """
    return theSystemPath.iterModules()

def getModule(moduleName):
    """
    Retrieve a module from the system path.
    """
    return theSystemPath[moduleName]


import numpy as np
import pycasso
from sklearn.model_selection import KFold
from sklearn.linear_model.coordinate_descent import _alpha_grid
from sklearn.model_selection import KFold
from sklearn.metrics import r2_score
from . import lm
import pdb

# Wrapper class to stitch together multiple path-wise solutions from
# pycasso corresponding to different gamma
class PycassoGrid():
	'''	class PycassoGrid : Fit the pycasso solver pathwise on a grid of regularization
		parameters.

		penalty : 'l1' 'scad' 'mcp' (for l1, can just use Pycassolasso if desired)

	'''

	def __init__(self, penalty, n_alphas=100, gamma = [3], fit_intercept = False,
				 eps = 1e-3, alphas=None):

		self.penalty = penalty
		self.n_alphas = n_alphas
		if np.isscalar(gamma):
			gamma = [gamma]
		self.gamma = np.array(gamma)
		self.fit_intercept = fit_intercept
		self.eps = eps
		self.alphas = alphas

	def fit(self, X, y):

		_, n_features = X.shape

		if self.alphas is None:
			self.alphas = self.get_alphas(X, y)
		else:
			self.n_alphas = self.alphas.size

		coefs = np.zeros((self.gamma.size, self.n_alphas, n_features))

		for gidx, gamma in enumerate(self.gamma):

			solver = pycasso.Solver(X, y, family = 'gaussian',
									penalty = self.penalty, gamma = gamma,
									useintercept = self.fit_intercept,
									lambdas = self.alphas)
			solver.train()
			coefs[gidx, ...] = solver.result['beta']

		self.coef_ = coefs

	def get_alphas(self, X, y):

		# The lambda selection that pycasso uses is essentially the same
		# as alpha_grid
		return _alpha_grid(X, y, n_alphas = self.n_alphas, eps = self.eps)

	# Predict over the entire grid
	def predict_grid(self, X):

		y_pred = self.coef_ @ X.T
		return y_pred

	# Score over the whole grid
	def score_grid(self, X, y):

		scores = np.zeros((self.gamma.size, self.n_alphas))

		y_pred = self.predict_grid(X)

		for gidx in range(self.gamma.size):

			scores[gidx, :] = np.array([r2_score(y, y_pred[gidx, j, :])
							  for j in range(self.n_alphas)])

		return scores

# Wrapper class to automate cross-validation with pycasso solvers
# penalty: 'l1' , 'mcp', or 'scad'
class PycassoCV(PycassoGrid):

	def __init__(self, penalty, n_alphas=100, gamma = [3],
				 nfolds = 5, fit_intercept = False, eps = 1e-3,
				 alphas=None):

		self.nfolds = nfolds
		super(PycassoCV, self).__init__(penalty, n_alphas, gamma,
										fit_intercept, eps, alphas)

	def fit(self, X, y):

		# The lambda selection that pycasso uses is essentially the same
		# as alpha_grid
		if self.alphas is None:
			self.alphas = _alpha_grid(X, y, n_alphas = self.n_alphas, eps = self.eps)
		else:
			self.n_alphas = self.alphas.size
		# Initialize cross-validator object
		self.cross_validator = KFold(n_splits = self.nfolds)

		# Store scores
		scores = np.zeros((self.nfolds, self.gamma.size, self.n_alphas))

		fold_idx = 0
		for train_idxs, test_idxs in self.cross_validator.split(X):
			X_train = X[train_idxs]
			y_train = y[train_idxs]

			X_test = X[test_idxs]
			y_test = y[test_idxs]

			super(PycassoCV, self).fit(X_train, y_train)

			scores[fold_idx, ...] = self.score_grid(X_test, y_test)

			fold_idx += 1

		# Average over folds
		scores = np.mean(scores ,axis = 0)

		self.scores = scores

		best_idx = np.unravel_index(np.argmax(scores.ravel()), scores.shape)

		# Set the selected parameters
		self.gamma_ = self.gamma[best_idx[0]]
		self.alpha_ = self.alphas[best_idx[1]]

		# Dummy regularization path (throw away all but the first)
		dummy_alphas = np.array([self.alpha_, self.alpha_/2, self.alpha_/4])

		# Refit with the selected parameters
		solver = pycasso.Solver(X, y, family = 'gaussian',
								penalty = self.penalty, gamma = self.gamma_,
								useintercept = self.fit_intercept,
								lambdas = dummy_alphas)

		solver.train()
		# Store final coefficients
		self.coef_ = solver.result['beta'][0, :]

# Cross-validator for PycassoElasticNet
class PycEnCV(lm.PycassoElasticNet):

	def __init__(self, n_folds=5, fit_intercept=False, max_iter=1000, lambda1 = None,
				 lambda2 = None):

		self.nfolds = n_folds
		super(PycEnCV, self).__init__(fit_intercept, max_iter, lambda1, lambda2)

	def fit(self, X, y):

		self.init_reg_params(X, y)
		if np.isscalar(self.lambda2):
			self.lambda2 = np.array([self.lambda2])

		cross_validator = KFold(n_splits = self.nfolds)

		scores = np.zeros((self.lambda2.size, self.lambda1.size,
						   self.nfolds))

		# Outer loop over lambda2
		for i, l2 in enumerate(self.lambda2):

			fold_idx = 0

			for train_idxs, test_idxs in cross_validator.split(X):
				X_train = X[train_idxs]
				y_train = y[train_idxs]

				X_test = X[test_idxs]
				y_test = y[test_idxs]

				en = lm.PycassoElasticNet(fit_intercept=self.fit_intercept, max_iter=self.max_iter,
									   lambda1 = self.lambda1, lambda2 = l2)
				en.fit(X_train, y_train)

				y_pred = en.coef_ @ X_test.T

				scores[i, :, fold_idx] = np.array([r2_score(y_test, y_pred[j, :]) for j in
												   range(self.lambda1.size)])
				fold_idx += 1

		# Average over folds
		scores = np.mean(scores, axis = -1)
		self.scores = scores
		best_idx = np.unravel_index(np.argmax(scores), (self.lambda2.size, self.lambda1.size))

		self.lambda2_ = self.lambda2[best_idx[0]]
		self.lambda1_ = self.lambda1[best_idx[1]]

		# Refit with the selected parameters.
		en = lm.PycassoElasticNet(fit_intercept = self.fit_intercept, max_iter=self.max_iter,
							   lambda1=self.lambda1_, lambda2=self.lambda2_)
		en.fit(X, y)

		self.coef_ = en.coef_


import numpy as np
import math
#from scipy.misc import imresize
import scipy.io
import matplotlib.pyplot as plt
from scipy import interpolate
from pytorch_models import train_SRCNN, train_DnCNN
import torch
from torch.utils.data import TensorDataset
from datetime import datetime
from tqdm import tqdm

##################################################################
#
##################################################################

n_epochs = 200
load_from_checkpoint = True

def psnr(target, ref):
    # assume RGB image
    target_data = np.array(target, dtype=float)
    ref_data = np.array(ref, dtype=float)

    diff = ref_data - target_data
    diff = diff.flatten('C')

    rmse = math.sqrt(np.mean(diff ** 2.))

    return 20 * math.log10(255. / rmse)

def interpolation(noisy , SNR , Number_of_pilot , interp, K, P, model_use="new"):
    N, N_S, N_D = noisy.shape
    noisy_image = np.zeros((N, N_S, N_D,2))

    noisy_image[:,:,:,0] = np.real(noisy)
    noisy_image[:,:,:,1] = np.imag(noisy)

    if model_use == "old":
        if (Number_of_pilot == 48):
            idx = [14*i for i in range(1, 72,6)]+[4+14*(i) for i in range(4, 72,6)]+[7+14*(i) for i in range(1, 72,6)]+[11+14*(i) for i in range(4, 72,6)]
        elif (Number_of_pilot == 16):
            idx= [4+14*(i) for i in range(1, 72,9)]+[9+14*(i) for i in range(4, 72,9)]
        elif (Number_of_pilot == 24):
            idx = [14*i for i in range(1,72,9)]+ [6+14*i for i in range(4,72,9)]+ [11+14*i for i in range(1,72,9)]
        elif (Number_of_pilot == 8):
            idx = [4+14*(i) for  i in range(5,72,18)]+[9+14*(i) for i in range(8,72,18)]
        elif (Number_of_pilot == 36):
            idx = [14*(i) for  i in range(1,72,6)]+[6+14*(i) for i in range(4,72,6)] + [11+14*i for i in range(1,72,6)]

        r = [x//14 for x in idx]
        c = [x%14 for x in idx]
    else:
        allCarriers = np.arange(K)  # indices of all subcarriers ([0, 1, ... K-1])

        pilotCarriers = allCarriers[::K//P]
        pilotCarriers = np.hstack([pilotCarriers, np.array([allCarriers[-1]])])
        r = []
        c= []
        temp = np.zeros(len(pilotCarriers), dtype=np.int)
        for i in range(14):
            r.append(pilotCarriers)
            c.append(temp+i)
        r = np.hstack(r)
        c = np.hstack(c)

    interp_noisy = np.zeros((N, N_S, N_D,2))

    for i in tqdm(range(len(noisy))):
        z = [noisy_image[i,j,k,0] for j,k in zip(r,c)]
        if(interp == 'rbf'):
            f = interpolate.Rbf(np.array(r).astype(float), np.array(c).astype(float), z,function='gaussian')
            X , Y = np.meshgrid(range(N_S),range(N_D))
            z_intp = f(X, Y)
            interp_noisy[i,:,:,0] = z_intp.T
        elif(interp == 'spline'):
            tck = interpolate.bisplrep(np.array(r).astype(float), np.array(c).astype(float), z)
            z_intp = interpolate.bisplev(range(N_S),range(N_D),tck)
            interp_noisy[i,:,:,0] = z_intp
        z = [noisy_image[i,j,k,1] for j,k in zip(r,c)]
        if(interp == 'rbf'):
            f = interpolate.Rbf(np.array(r).astype(float), np.array(c).astype(float), z,function='gaussian')
            X , Y = np.meshgrid(range(N_S),range(N_D))
            z_intp = f(X, Y)
            interp_noisy[i,:,:,1] = z_intp.T
        elif(interp == 'spline'):
            tck = interpolate.bisplrep(np.array(r).astype(float), np.array(c).astype(float), z)
            z_intp = interpolate.bisplev(range(N_S),range(N_D),tck)
            interp_noisy[i,:,:,1] = z_intp

    interp_noisy = np.concatenate((interp_noisy[:,:,:,0], interp_noisy[:,:,:,1]), axis=0).reshape(2*N, N_S, N_D, 1)

    return interp_noisy

if __name__ == "__main__":
    # load datasets
    channel_model = "VehA"
    SNR = 25
    Number_of_pilots = 48
    Number_of_Subcarriers = 64
    Number_of_Pilots_per_symbol = 8 # Comb type OFDM

    print("Reading the Noisy and Perfect Data files in .mat format.")

    # perfect = scipy.io.loadmat("data\Perfect_H_40000.mat") ['My_perfect_H']
    # noisy_input = scipy.io.loadmat("data\My_Noisy_H_12.mat") ['My_noisy_H']

    perfect = scipy.io.loadmat("data\\Perfect_H_dataset_for_" + str(SNR) + "db") ['Perfect_H']
    noisy_input = scipy.io.loadmat("data\\Noisy_H_dataset_for_" + str(SNR) + "db") ['Noisy_H']

    print(" \n ")
    print("Read the Noisy and Perfect Data files.")
    print(" \n ")
    print("Interpolating the noisy data into LR images")
    print(f"SNR: {SNR}\nNumber of pilots: {Number_of_pilots}")
    interp_noisy = interpolation(noisy_input , SNR , Number_of_pilots , 'rbf', K=Number_of_Subcarriers, P=Number_of_Pilots_per_symbol, model_use="new")

    N, N_S, N_D = perfect.shape
    perfect_image = np.zeros((N, N_S, N_D, 2))
    perfect_image[:,:,:,0] = np.real(perfect)
    perfect_image[:,:,:,1] = np.imag(perfect)
    perfect_image = np.concatenate((perfect_image[:,:,:,0], perfect_image[:,:,:,1]), axis=0).reshape(2*N, N_S, N_D, 1)

    import pdb;pdb.set_trace()
    print("==========================================================")

    ####### ------ training SRCNN ------ #######
    interp_noisy = interp_noisy.transpose((0,3,1,2))
    perfect_image = perfect_image.transpose((0,3,1,2))

    t_interp_noisy = torch.Tensor(interp_noisy)
    t_perfect_image = torch.Tensor(perfect_image)
    my_dataset = TensorDataset(t_interp_noisy, t_perfect_image)

    train_size = int(0.8 * len(my_dataset))
    test_size = len(my_dataset) - train_size

    device=torch.device("cuda:0" if torch.cuda.is_available else "cpu")
    model_name = "SRCNN"

    print("Data is ready for training.\nTraining the SRCNN first....")

    train_SRCNN(train_size, test_size, dataset=my_dataset, n_epochs=1000, model_name=model_name, device=device, load_from_checkpoint=load_from_checkpoint)

    print("Data is ready for training.\nTraining the DnCNN first....")
    model_name = "DnCNN"

    train_DnCNN(train_size, test_size, dataset=my_dataset, n_epochs=n_epochs, model_name=model_name, device=device, load_from_checkpoint=load_from_checkpoint, path_to_SRCNN="saved_models\SRCNN_checkpoint_latest.pt")

    print("Trained weights stored in 'saved_models' folder.")


# helper script to convert CSV files into a single SQLite database file

import sqlalchemy as sa
import pandas as pd
from pathlib import Path

def create_table_sql(df, tbl_name):
    cols = []
    pks = []
    fks = []
    for c in df.columns:
        if c == f'{tbl_name}_id':
            pks.append(f'PRIMARY KEY("{c}")')
        elif c.endswith('_id'):
            fks.append(f'FOREIGN KEY("{c}") REFERENCES "{c[:-3]}"("{c}")')
        dtype = df[c].dtype
        if pd.api.types.is_integer_dtype(dtype):
            ctype = 'BIGINT'
        elif pd.api.types.is_float_dtype(dtype):
            ctype = 'FLOAT'
        elif pd.api.types.is_datetime64_any_dtype(dtype):
            ctype = 'DATETIME'
        else:
            ctype = 'TEXT'
        cols.append(f'"{c}" {ctype}')
    stmt = f'CREATE TABLE "{tbl_name}" ({", ".join(cols + pks + fks)})'
    return stmt

engine = sa.create_engine('sqlite+pysqlite:///berka-sqlite.db', echo=False)

data = {}
for fn in Path('.').glob('*.csv'):
    df = pd.read_csv(fn)
    # convert dtypes
    for col in df.columns:
        if col in ['date', 'issued']:
            df[col] = pd.to_datetime(df[col])
        if col.endswith('_id'):
            df[col] = df[col].astype(str)
    # get filename w/o extension
    tbl_name = fn.stem
    data[tbl_name] = df

with engine.connect() as conn:
    for tbl_name, df in data.items():
        # create table
        stmt = create_table_sql(df, tbl_name)
        conn.execute(sa.text(stmt))
        print(f"created table {tbl_name}")
    conn.commit()
    conn.close()

with engine.connect() as conn:
    for tbl_name, df in data.items():
        # insert records
        df.to_sql(tbl_name, conn, index=False, if_exists='append')
        print(f"loaded data to {tbl_name}")
    conn.commit()
    conn.close()