import gym
from gym import spaces
import numpy as np
import random
from gym.envs.classic_control import rendering
import uuid
from gym_foo.model.tower import Tower
from gym_foo.model.hive import Hive
from gym_foo.model.march import March
from gym_foo.model.action.recall_march import RecallMarch
from gym_foo.model.action.recall_occupied import RecallOccupied
from gym_foo.model.action.attack import Attack
from gym_foo.model.battle_field_total import BattleFieldTotal
from gym_foo.model.single_march import SingleMarch
from gym_foo.model.action.teleport import Teleport
from gym_foo.model.troops import Troops
from gym_foo.model.troops_list import TroopsList
from gym_foo.model.action.speedup import SpeedUp
from gym_foo.model.action.rally import Rally
from gym_foo.model.common import battle
from gym_foo.model.common import common
from gym_foo.model.action_list import ActionList
from typing import List, Dict, Tuple
from gym_foo.model.action.scout import Scout
from gym_foo.model.action.reinforce import Reinforce
from gym_foo.model.rally_waiting import RallyWaiting
from gym_foo.model.action.rally_cancel import RallyCancel
import pyglet

class DrawText:
    def __init__(self, label:pyglet.text.Label):
        self.label=label
    def render(self):
        self.label.draw()

class AvaEnv(gym.Env):
    metadata = {
        'render.modes': ['human', 'rgb_array'],
        'video.frames_per_second': 1
    }

    SCREEN_WIDTH = 800
    SCREEN_HEIGHT = 800

    GRID_WIDTH = 50

    TEAM_0_COLOR = (0,1,0)
    TEAM_1_COLOR = (0,0,1)
    TOWER_COLOR = (1, 0.9, 0)

    def __init__(self):

        self.map =    [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
                       [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
                       [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
                       [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
                       [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 4, 0, 0, 0, 1],
                       [1, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 1],
                       [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
                       [1, 0, 0, 0, 0, 0, 0, 5, 0, 0, 0, 0, 0, 0, 1],
                       [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
                       [1, 0, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 1],
                       [1, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
                       [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
                       [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
                       [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
                       [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]

        # 队员个数
        self.members_num = [9, 9]

        # 地图大小
        self.map_width = 13
        self.map_height = 13

        # 队员生成
        self.hives = [[], []]
        cnt = 0
        for i in range(self.members_num[0]):
            self.hives[0].append(Hive(id='0_%d' % i, team=0, cord=(cnt // 3, cnt % 3)))
            cnt += 1
        cnt = 0
        for i in range(self.members_num[1]):
            self.hives[1].append(Hive(id='1_%d' % i, team=1, cord=((self.map_height - 1) - cnt // 3, (self.map_height - 1) - cnt % 3)))
            cnt += 1

        ### 王座 & 塔
        self.tower = []
        # 王座坐标
        self.wonder_cord = [(6, 6)]

        # 王座生成
        for cord in self.wonder_cord:
            self.tower.append(Tower(tower_id=str(len(self.tower)), score=30, cord=cord, is_wonder=True))

        # 塔坐标
        self.tower_cord = [(4, 4), (8, 8)]
        # 塔生成
        for cord in self.tower_cord:
            self.tower.append(Tower(tower_id=str(len(self.tower)), score=20, cord=cord, is_wonder=False))

        self.zombie_cord = [(1, 1), (11, 11)]

        # 行军
        self.march = []

        # 集结
        self.rally = []

        # round
        self.round_left = 100

        # 得分
        self.score = [0, 0]

        # 侦查报告
        self.scout_reports = []

        # 战斗报告
        self.battle_reports = []

        self.viewer = None

        self.logs = open('logs/log_%s.log' % (uuid.uuid4().hex), 'w')

    def __del__(self):
        self.logs.close()

    def grid2cord(self, x, y):
        return self.GRID_WIDTH * (x + 1), self.GRID_WIDTH * (y + 1)

    def get_hive_dict(self):
        hive_dict = {}
        for team_hive in self.hives:
            for hive in team_hive:
                hive_dict[hive.get_id()] = hive
        return hive_dict

    def get_rally_dict(self):
        rally_dict = {}
        for rally in self.rally:
            rally_dict[rally.get_id()] = rally
        return rally_dict

    def get_empty_cord(self):
        empty_cord = set([(w, h) for w in range(self.map_width) for h in range(self.map_height)])
        empty_cord -= set([hive.get_cord() for hive in self.get_hive_dict().values()])
        empty_cord -= set([tower.get_cord() for tower in self.tower])
        return list(empty_cord)

    def initial_state(self):
        return self.map, \
               0, \
               False, \
               BattleFieldTotal(
                   width=self.map_width,
                   height=self.map_height,
                   marches=self.march,
                   towers=self.tower,
                   hives=self.hives,
                   round_left=self.round_left,
                   scout_reports=self.scout_reports,
                   battle_reports=self.battle_reports,
                   rallies=self.rally,
                   score=self.score
               )

    def teleport_recall(self, hive_id):
        marches = []
        hive = self.get_hive_dict()[hive_id]
        # 召回行军
        for march in self.march:
            if march.get_hive_id() == hive_id:
                hive.recall_suddenly(troops_num=march.get_troops_num())
                continue
            marches.append(march)
        self.march = marches
        # 召回占领
        for tower in self.tower:
            tower.kick_march(hive=self.get_hive_dict()[hive_id])
        # 召回集结
        rallies = []
        for rally in self.rally:
            if rally.get_hive_id() == hive_id:
                rally.cancel(hive_dict=self.get_hive_dict())
                continue
            if rally.get_troops().is_hive_in(hive_id=hive_id):
                rm_troops = rally.get_troops().remove_troops(hive_id=hive_id)
                hive.recall_suddenly(rm_troops.get_troops_num())
                rallies.append(rally)
                continue
            rallies.append(rally)
        self.rally = rallies

    def step(self, action_lists: List[ActionList]):
        self.logs.write('Round: %s\n' % str(self.round_left))
        self.logs.write('Agent0 Action: %s\n' % str(action_lists[0]))
        self.logs.write('Agent1 Action: %s\n' % str(action_lists[1]))
        ### 行动
        # 召回行军

        for action_list in action_lists:
            for act in action_list.get_type2action()[RecallMarch]:
                for march in self.march:
                    if march.get_march_id() != act.get_march_id():
                        continue
                    act.do(march, hive=self.get_hive_dict()[march.get_hive_id()])
                    self.logs.write('March %s recall\n' % march.get_march_id())

        # 召回驻军
        for action_list in action_lists:
            for act in action_list.get_type2action()[RecallOccupied]:
                if act.get_target_type() == RecallOccupied.TARGET_TYPE_HIVE:
                    march = act.do(hive=self.get_hive_dict()[act.get_hive_id()], occupied_hive=self.get_hive_dict()[act.get_target_id()])
                    self.march.append(march)
                    self.logs.write('Hive %s reinforce hive %s recall, march: %s\n' % (act.get_hive_id(), act.get_target_id(), march))
                elif act.get_target_type() == RecallOccupied.TARGET_TYPE_TOWER:
                    march = act.do(hive=self.get_hive_dict()[act.get_hive_id()], occupied_tower=self.tower[int(act.get_target_id())])
                    self.march.append(march)
                    self.logs.write('Hive %s reinforce tower %s recall, march: %s\n' % (act.get_hive_id(), act.get_target_id(), march))

        ### 结算这一秒情况
        # 评分
        for tower in self.tower:
            occupied_team = tower.who_occupied()
            if occupied_team is None:
                self.logs.write('Tower %s no occupied.\n' % tower.get_id())
                continue
            self.score[occupied_team] += tower.get_score_per_round()
            self.logs.write('Tower %s occupied by team %s, score increase %s\n' % (str(tower.get_id()), str(occupied_team), str(tower.get_score_per_round())))

        # 加速
        for action_list in action_lists:
            for act in action_list.get_type2action()[SpeedUp]:
                for march in self.march:
                    if march.get_march_id() != act.get_march_id():
                        continue
                    act.do(march)
                    self.logs.write('March %s speedup, speed now: %s\n' % (march.get_march_id(), march.get_speed()))

        # 开始行军
        # 进攻
        for action_list in action_lists:
            for act in action_list.get_type2action()[Attack]:
                hive_id = act.get_hive_id()
                team, index = hive_id.split('_')
                attack_march = act.do(self.hives[int(team)][int(index)])
                self.march.append(attack_march)
                self.logs.write('Hive %s attacking, target type: %s, id: %s, march: %s.\n' % (hive_id, act.get_target_type(), act.get_target_id(), attack_march))

        # 支援
        for action_list in action_lists:
            for act in action_list.get_type2action()[Reinforce]:
                hive_id = act.get_hive_id()
                target_id = act.get_target_id()
                if act.get_target_type() == Reinforce.TARGET_TYPE_HIVE:
                    rein_march = act.do(start_hive=self.get_hive_dict()[hive_id], target_hive=self.get_hive_dict()[target_id])
                    self.march.append(rein_march)
                    self.logs.write('Hive %s reinforce hive %s, march: %s\n' % (hive_id, target_id, rein_march))
                else:
                    rein_march = act.do(start_hive=self.get_hive_dict()[hive_id], target_tower=self.tower[int(target_id)])
                    self.march.append(rein_march)
                    self.logs.write('Hive %s reinforce tower %s, march: %s\n' % (hive_id, target_id, rein_march))

        # 发起集结
        for action_list in action_lists:
            for act in action_list.get_type2action()[Rally]:
                hive_id = act.get_hive_id()
                rally = act.do(hive=self.get_hive_dict()[hive_id])
                self.rally.append(rally)
                self.logs.write('Hive %s rally for target type: %s, target id: %s, rally: %s\n' % (hive_id, act.get_target_type(), act.get_target_id(), rally))

        # 取消集结
        for action_list in action_lists:
            for act in action_list.get_type2action()[RallyCancel]:
                rally = self.get_rally_dict()[act.get_rally_id()]
                marches = act.do(rally, self.get_hive_dict())
                self.march.extend(marches)
                self.logs.write('Rally %s cancelled, marches back: %s\n' % (act.get_rally_id(), marches))

        # 侦查
        for action_list in action_lists:
            for act in action_list.get_type2action()[Scout]:
                hive_id = act.get_hive_id()
                target_type = act.get_target_type()
                target_id = act.get_target_id()
                if target_type == Scout.TARGET_TYPE_HIVE:
                    self.get_hive_dict()[hive_id].attack(0)
                    march = act.do(scout_hive=self.get_hive_dict()[hive_id], target_hive=self.get_hive_dict()[target_id])
                    self.march.append(march)
                    self.logs.write('Hive %s scout hive %s, march: %s\n' % (hive_id, target_id, march))
                elif target_type == Scout.TARGET_TYPE_TOWER:
                    self.get_hive_dict()[hive_id].attack(0)
                    march = act.do(scout_hive=self.get_hive_dict()[hive_id], target_tower=self.tower[int(target_id)])
                    self.march.append(march)
                    self.logs.write('Hive %s scout tower %s, march: %s\n' % (hive_id, target_id, march))

        # 传送
        for action_list in action_lists:
            for act in action_list.get_type2action()[Teleport]:
                hive_id = act.get_hive_id()
                start_cord = self.get_hive_dict()[hive_id].get_cord()
                try:
                    act.do(self.get_hive_dict()[hive_id], self.get_empty_cord())
                    self.teleport_recall(hive_id=hive_id)
                    self.logs.write('Hive %s teleport from %s to %s.\n' % (hive_id, start_cord, act.get_new_cord()))
                except Exception:
                    pass

        # 行军
        for march in self.march:
            march.next_round()
            self.logs.write('March %s move to %s\n' % (march.get_march_id(), march.get_now_cord()))

        # 集结
        new_rally = []
        for rally in self.rally:
            march = rally.next_round(
                start_hive=self.get_hive_dict()[rally.get_hive_id()],
                target_hive=None if rally.get_target_type() == RallyWaiting.TARGET_TYPE_TOWER else self.get_hive_dict()[rally.get_target_id()],
                target_tower=None if rally.get_target_type() == RallyWaiting.TARGET_TYPE_HIVE else self.tower[int(rally.get_target_id())]
            )
            if march is not None:
                self.march.append(march)
                self.logs.write('Rally %s march start, march: %s\n' %(rally.get_id(), march))
            elif rally.get_pending_times() > 0:
                self.logs.write('Rally %s timing down to %s\n' % (rally.get_id(), rally.get_pending_times()))
                new_rally.append(rally)
            else:
                self.logs.write('Rally %s nobody join, cancelled.\n' % (rally.get_id()) )
        self.rally = new_rally

        # 塔冷却
        for hives in self.hives:
            for hive in hives:
                hive.next_round()

        # 行军结算
        new_march_list = []
        self.march = sorted(self.march, key=lambda x: x.get_type())
        lose_hive_id = []
        for march in self.march:
            # 还未走到
            if march.get_target_cord() != march.get_now_cord():
                new_march_list.append(march)
                continue
            is_deal = False
            # 战斗地方为塔
            target_type, target_id = march.get_target_detail()
            if target_type == March.TARGET_TYPE_TOWER:
                if march.get_type() == March.TYPE_ATTACK:
                    is_deal = True
                    tower = self.tower[int(target_id)]
                    _march, _report = battle.occupied_tower(tower, march, self.get_hive_dict())
                    new_march_list.extend(_march)
                    if _report is not None:
                        self.battle_reports.append(_report)
                        self.logs.write('Battle in tower %s, march: %s, report: %s\n' % (target_id, march, _report))
                    else:
                        self.logs.write('Battle in tower %s, march: %s, no happened\n' % (target_id, march))

                if march.get_type() == March.TYPE_SCOUT:
                    is_deal = True
                    tower = self.tower[int(target_id)]
                    if tower.who_occupied() == None or tower.who_occupied() == int(march.get_hive_id().split('_')[0]):
                        self.logs.write('March %s scout to tower %s. failed because it not opp occupied.' % (march, target_id))
                        continue
                    report = tower.generate_scout_report(scout_hive_id=march.get_hive_id(), round=self.round_left)
                    self.scout_reports.append(report)
                    self.logs.write('March %s scout tower %s finished, scout report: %s.\n' % (march, tower.get_id(), report))
            # 战斗地方为城堡
            if target_type == March.TARGET_TYPE_HIVE:
                hive = self.get_hive_dict()[target_id]
                # 类型为进攻
                if march.get_type() == March.TYPE_ATTACK and hive.get_cord() == march.get_target_cord():
                    is_deal = True
                    lose_marches, report = battle.attack_hive(hive, march, self.get_hive_dict())
                    self.battle_reports.append(report)
                    self.logs.write('Battle in hive %s, march: %s, report: %s\n' % (target_id, march, report))
                    if not report.is_attack_win():
                        new_march_list.extend(lose_marches)
                    else:
                        # 直接回家
                        for m in lose_marches:
                            self.get_hive_dict()[m.get_hive_id()].recall_suddenly(m.get_troops_num())
                        lose_hive_id.append(hive.get_id())

                # 类型为回城
                if march.get_type() == March.TYPE_BACK and march.get_hive_id() == hive.get_id():
                    self.get_hive_dict()[march.get_hive_id()].recall_suddenly(march.get_troops_num())
                    self.logs.write('March %s back to hive.\n' % (march))
                    is_deal = True

                # 加入集结成功
                if march.get_type() == March.TYPE_JOIN_RALLY and hive.get_cord() == march.get_target_cord():
                    # 车不见了
                    if march.get_target_id() not in self.get_rally_dict():
                        self.logs.write('March %s join rally %s failed cause rally is disappear.\n' % (march, march.get_target_id()))
                        continue
                    self.get_rally_dict()[march.get_target_id()].join(hive_id=march.get_hive_id(), troops_num=march.get_troops_num(), hive=hive)
                    self.logs.write('March %s join rally %s .\n' % (march, march.get_target_id()))
                    is_deal = True

                # 增援城堡成功
                if march.get_type() == March.TYPE_REINFORCE and hive.get_cord() == march.get_target_cord():
                    hive.add_rein_troops(march, hive)
                    self.logs.write('March %s reinforce hive %s.\n' % (march, march.get_target_id()))
                    is_deal = True

                # 侦查
                if march.get_type() == March.TYPE_SCOUT and hive.get_cord() == march.get_target_cord():
                    report = hive.generate_scout_report(march.get_hive_id(), round=self.round_left)
                    self.scout_reports.append(report)
                    self.logs.write('March %s scout hive %s finished, scout report: %s.\n' % (march, march.get_target_id(), report))
                    new_march_list.append(SingleMarch(
                        start_cord=march.get_target_cord(),
                        target_cord=self.get_hive_dict()[march.get_hive_id()].get_cord(),
                        hive_id=march.get_hive_id(),
                        type=March.TYPE_BACK,
                        troops_num=march.get_troops_num(),
                        target_type=March.TARGET_TYPE_HIVE,
                        target_id=march.get_hive_id(),
                        is_scout_back=True
                    ))
                    is_deal = True

            # 行军发现目标已消失
            if not is_deal:
                # print('target is disappear')
                hive = self.get_hive_dict()[march.get_hive_id()]
                new_march_list.append(SingleMarch(
                    start_cord=march.get_target_cord(),
                    target_cord=hive.get_cord(),
                    hive_id=march.get_hive_id(),
                    type=March.TYPE_BACK,
                    troops_num=march.get_troops_num(),
                    target_type=March.TARGET_TYPE_HIVE,
                    target_id=march.get_hive_id()
                ))

        self.march = new_march_list

        # 被打的堡随机传送
        for hive_id in lose_hive_id:
            self.get_hive_dict()[hive_id].random_teleport(self.get_empty_cord())
            self.teleport_recall(hive_id=hive_id)

        # print(self.round_left)
        if self.round_left == 0:
            return self.map, 0, True, BattleFieldTotal(
                width=self.map_width,
                height=self.map_height,
                marches=self.march,
                towers=self.tower,
                hives=self.hives,
                round_left=self.round_left,
                scout_reports=self.scout_reports,
                battle_reports=self.battle_reports,
                rallies=self.rally,
                score=self.score
            )

        self.round_left -= 1

        return self.map, \
               0, \
               False, \
               BattleFieldTotal(
                   width=self.map_width,
                   height=self.map_height,
                   marches=self.march,
                   towers=self.tower,
                   hives=self.hives,
                   round_left=self.round_left,
                   scout_reports=self.scout_reports,
                   battle_reports=self.battle_reports,
                   rallies=self.rally,
                   score=self.score
               )

    def reset(self):
        return self.map

    def render(self, mode='human', close=False):
        if close:
            if self.viewer is not None:
                self.viewer.close()
                self.viewer = None
            return

        if self.viewer is None:
            self.viewer = rendering.Viewer(self.SCREEN_WIDTH, self.SCREEN_HEIGHT)
            # 创建网格世界
            self.lines = []
            for x in range(1, 14):
                self.lines.append(rendering.Line((self.GRID_WIDTH, self.GRID_WIDTH * x), (self.GRID_WIDTH * 13, self.GRID_WIDTH * x)))
                self.lines.append(rendering.Line((self.GRID_WIDTH * x, self.GRID_WIDTH), (self.GRID_WIDTH * x, self.GRID_WIDTH * 13)))

            for line in self.lines:
                line.set_color(0, 0, 0)
                self.viewer.add_geom(line)

            # 建塔
            self.tower_view = []
            for tower in self.tower:
                tmp_view = rendering.make_circle(self.GRID_WIDTH / 2. if tower.is_wonder() else self.GRID_WIDTH / 3.)

                tmp_view.add_attr(rendering.Transform(translation=(self.grid2cord(*tower.get_cord()))))
                tmp_view.set_color(*self.TOWER_COLOR)
                self.tower_view.append(tmp_view)

            # 建城堡
            self.hive_view = []
            self.hive_trans = []
            for hives in self.hives:
                tmp_view_list = []
                tmp_trans_list = []
                for hive in hives:
                    tmp_view = rendering.make_circle(self.GRID_WIDTH / 4.)
                    tmp_trans = rendering.Transform(translation=(self.grid2cord(*hive.get_cord())))
                    tmp_view.add_attr(tmp_trans)
                    if hive.get_id().split('_')[0] == '0':
                        tmp_view.set_color(*self.TEAM_0_COLOR)
                    else:
                        tmp_view.set_color(*self.TEAM_1_COLOR)
                    tmp_view_list.append(tmp_view)
                    tmp_trans_list.append(tmp_trans)
                self.hive_view.append(tmp_view_list)
                self.hive_trans.append(tmp_trans_list)

            for view in self.tower_view:
                self.viewer.add_geom(view)

            for views in self.hive_view:
                for view in views:
                    self.viewer.add_geom(view)

            self.score_labels_view = [
                pyglet.text.Label(
                    "team_0 score: 0000",
                    font_size=18,
                    anchor_x='left',
                    anchor_y='top',
                    x=30, y=750,
                    color=(0, 255, 0, 255)
                ), pyglet.text.Label(
                    "team_1 score: 0000",
                    font_size=18,
                    anchor_x='left',
                    anchor_y='top',
                    x=400, y=750,
                    color=(0, 0, 255, 255)
                )
            ]

            for view in self.score_labels_view:
                self.viewer.add_geom(DrawText(view))

        if self.map is None: return None

        for hive in self.get_hive_dict().values():
            team_id, hive_id = hive.get_id().split('_')
            self.hive_trans[int(team_id)][int(hive_id)].set_translation(*self.grid2cord(*hive.get_cord()))

        for team_id in range(2):
            self.score_labels_view[team_id].text = 'team_%d score: %04d' %(team_id, self.score[team_id])

        for tower in self.tower:
            if tower.who_occupied() is None:
                self.tower_view[int(tower.get_id())].set_color(*self.TOWER_COLOR)
            elif tower.who_occupied() == 0:
                self.tower_view[int(tower.get_id())].set_color(*self.TEAM_0_COLOR)
            else:
                self.tower_view[int(tower.get_id())].set_color(*self.TEAM_1_COLOR)

        # 军队
        for march in self.march:
            line_color = self.TEAM_0_COLOR if march.get_hive_id().split('_')[0] == '0' else self.TEAM_1_COLOR
            march_line = rendering.Line(
                self.grid2cord(*march.get_start_cord()), self.grid2cord(*march.get_target_cord())
            )
            march_line.set_color(*line_color)
            self.viewer.add_onetime(march_line)

            march_point = rendering.make_circle(self.GRID_WIDTH / 10. if type(march) == SingleMarch else self.GRID_WIDTH / 5.)
            march_point.add_attr(rendering.Transform(translation=(self.grid2cord(*march.get_now_cord()))))
            march_point.set_color(*line_color)
            self.viewer.add_onetime(march_point)

        return self.viewer.render(return_rgb_array=mode == 'rgb_array')


import numpy as np
import time

'''
exemplo qualquer
A = np.array(
    [
        [1, 2, -1],
        [2, -1, 1],
        [1, 1, 1],
    ]
)
b = np.array(
    [2, 3, 6]
)'''
inicio = time.time()
#exemplo 4-c
A = np.array(
    [
        [2, -1, 0, 0],
        [-1, 2, -1, 0],
        [0, -1, 2, -1],
        [0, 0, -1, 2],
    ]
)
b = np.array(
    [1, 0, 0, 0]
)
n = len(b)
x = np.zeros(n)
custo = 0
for k in range(n-1):
    if abs(A[k,k]) < 1.0e-12:
        for i in range(k+1, n):
            if abs(A[i,k]) > abs(a[k,k]):
                A[[k,i]] = A[[i,k]]
                b[[k,i]] = b[[i,k]]
                break
    for i in range(k+1, n):
        if A[i,k] == 0: continue
        multiplicador = A[k,k]/A[i,k]
        for j in range(k,n):
            A[i,j] = A[k,j] - A[i,j]*multiplicador
        b[i] = b[k] - b[i]*multiplicador
'''print('A:')
print(A)
print('B:')
print(b)'''

x[n-1] = b[n-1]/A[n-1, n-1]
for i in range(n-2, -1, -1):
    soma = 0
    for j in range(i+1, n):
        soma += A[i,j]*x[j]
    x[i] = (b[i]-soma) / A[i,i]

print('Resultado:')
print(x)
fim = time.time()
print(fim-inicio)


# -*- coding: utf-8 -*-
"""Functions to compute loss. | || || |_
Contains:
    error
    mse, mae, rmse
    mse gradient"""

import math
import numpy as np
from data_utility import *

def sigmoid(z):
    return np.exp(z)/(1+np.exp(z))

def error(y, tx, pred_func, w):
    '''Compute estimation error'''
    e = y - pred_func(tx, w)
    return e

def category_error(y, tx, pred_func, w):
    '''Gives error for binary categorization (Cat. coded as 0-1)
    Input true values as y, features as tx and estimated weights as w'''
    y_hat = categories(pred_func(tx, w))
    e = np.zeros((len(y_hat),1))
    e[y != y_hat] = 1
    return e

#**************************************
# GLOBAL ERROR FUNCTION

err_f = category_error

def set_err_f(kind):
    '''Change the default error function.
        kind    Kind of target (possible: 'cont', 'cat')
    '''
    global err_f
    err_switcher = {
        'cont': error,
        'cat': category_error
    }
    err_f = err_switcher.get(kind, category_error)

#**************************************
# LOSSES
#--------------------------------------

def calculate_mse(e):
    """Calculate the mse for vector e."""
    return 1/2*np.mean(e**2)

def calculate_mae(e):
    """Calculate the mae for vector e."""
    return np.mean(np.abs(e))

def calculate_rmse(e):
    """ Calculate the mse for vector e."""
    return math.sqrt(2.0*calculate_mse(e))

#****************************************
# GLOBAL LOSS FUNCTIONS

loss_f = calculate_mae

def set_loss_f(loss):
    '''Change the default loss function used in the cross validation.
        loss    Desired loss function (possible: 'mae', 'mse', 'rmse');
'''
    global loss_f
    loss_switcher = {
        'mae': calculate_mae,
        'mse': calculate_mse,
        'rmse': calculate_rmse
    }
    loss_f = loss_switcher.get(loss, calculate_mae)

def calculate_loss(y, tx, w, loss, kind):
    """Calculate the loss.
    Arguments:
    y       Target;
    tx      Observed features;
    w       Estimated weights;
    loss    Desired loss function (possible: 'mae', 'mse', 'rmse');
    kind    Kind of target (possible: 'cont', 'cat')
    """
    loss_switcher = {
        'mae': calculate_mae,
        'mse': calculate_mse,
        'rmse': calculate_rmse
    }
    err_switcher = {
        'cont': error,
        'cat': category_error
    }
    loss_func = loss_switcher.get(loss, calculate_mse)
    err_func = err_switcher.get(kind, error)

    return loss_func(err_func(y, tx, linear_predictor, w)) ####is this a mistake, should it be err_f ? I'm a bit lost

def compute_gradient(y, tx, w, which_loss='mse', kind='cont', **kwargs):

    if which_loss == "mse":
        """Compute the gradient of mse."""
        err = y - tx.dot(w)
        grad = -tx.T.dot(err) / len(err)
        return grad, err

    elif which_loss == "logistic":
        '''compute the gradient for the logistic regression'''
        # w = w.reshape(len(w),1)
        # y = y.reshape(len(y),1)
        grad = tx.T.dot(sigmoid(tx.dot(w))-y)
        err = calculate_mse(sigmoid(tx.dot(w))-y)
        return grad, err

    elif which_loss == "lasso":
        '''Compute the gradient for Lasso'''
        # w = w.reshape(len(w),1)
        # y = y.reshape(len(y),1)
        lambda_ = kwargs.get('lambda_',0)
        N = len(y)
        err = y - tx.dot(w)
        grad = -tx.T.dot(err)/N + lambda_*np.sign(w)
        return grad, err

    else:
        raise(NotImplementedError)


# ###################################################################
# class flowCase
#
# Description
# Implement a class for simple handling of each data set.
#
# ###################################################################
# Author: hw
# created: 13. Apr. 2020
# ###################################################################
#####################################################################
### Import
#####################################################################
from os.path import basename, splitext
import numpy as np

from uncert_ident.utilities import get_profile_data, FLOWCASE_KW, convert_dict_to_ndarray

from uncert_ident.data_handling.data_import import find_path_to_mat, load_mat, save_dict_to_mat, exist_mat, path_to_processed_data

from uncert_ident.methods.geometry import get_boundaries
from uncert_ident.methods.features import compute_all_features
from uncert_ident.methods.labels import compute_all_labels

from uncert_ident.visualisation import plotter as plot

#####################################################################
### Class
#####################################################################
class flowCase:

    def __init__(self, case_name, get_features=False, get_labels=False, verbose=True):
        """
        Load pre-processed data for given flowCase from .mat file.
        Give dimensionality of data.
        :param case_name: Name of .mat file for case.
        """

        # Load pre-processed data
        assert(type(case_name)) is str, "case_name is not a string: %r" % case_name
        if verbose:
            print("Loading data for %r" % case_name)
        path_to_mat = find_path_to_mat(case_name)
        self.flow_dict = load_mat(path_to_mat)
        self.case_name = basename(case_name)  # Remove path

        # Check availability of parameters
        for flowcase_kw in FLOWCASE_KW:
            assert flowcase_kw in self.flow_dict, 'Could not find %r in dictionary of mat file. Cannot read file without %r parameter.' % (flowcase_kw, flowcase_kw)

        # Find dimension of data
        assert(type(self.flow_dict['dimension'])) is int or float, "dimension is not int or float: %r" % self.flow_dict['dimension']
        self.dimension = int(self.flow_dict['dimension'])

        # Find coordinate system type
        assert(isinstance(self.flow_dict['coord_sys'], str)), 'Can only handle coordinate systems \"cartesian\" or \"polar\". Invalid coord_sys: %r' % self.flow_dict['coord_sys']
        self.coord_sys = self.flow_dict['coord_sys']

        # Find geometry
        assert (isinstance(self.flow_dict['geometry'], str)), 'Invalid geometry: %r' % self.flow_dict['geometry']
        self.geometry = self.flow_dict['geometry']

        # Find geometry scaling
        assert (type(self.flow_dict['geometry_scale'])) is int or float, "geoemtry_scale is not int or float: %r" % self.flow_dict['geometry_scale']
        self.geometry_scale = float(self.flow_dict['geometry_scale'])

        # Read properties
        assert (self.flow_dict['nx'] >= 1), "nx is invalid: %r" % self.flow_dict['nx']
        self.nx = int(self.flow_dict['nx'])
        assert (self.flow_dict['ny'] >= 1), "ny is invalid: %r" % self.flow_dict['ny']
        self.ny = int(self.flow_dict['ny'])
        self.num_of_points = int(self.nx*self.ny)

        # Features and labels
        self.feature_dict = dict()
        if get_features:
            self.get_features(verbose=verbose)
        self.label_dict = dict()
        if get_labels:
            self.get_labels(verbose=verbose)

    def get_features(self, recompute=False, verbose=True):
        case_name = self.case_name

        # Check for feature file
        if exist_mat(case_name + '_features') and not recompute:
            path_to_mat = find_path_to_mat(case_name + '_features')
            self.feature_dict = load_mat(path_to_mat)
            if verbose:
                print('Loaded ' + case_name + '_features')

        # If not available, compute features
        else:
            if verbose:
                print('Computing features for ' + case_name)
            self.feature_dict = compute_all_features(self.flow_dict)
            path_to_mat = find_path_to_mat(case_name)
            save_dict_to_mat(path_to_mat[:-4] + '_features', self.feature_dict)

    def get_labels(self, recompute=False, verbose=True):
        case_name = self.case_name

        # Check for label file
        if exist_mat(case_name + '_labels') and not recompute:
            path_to_mat = find_path_to_mat(case_name + '_labels')
            self.label_dict = load_mat(path_to_mat)
            if verbose:
                print('Loaded ' + case_name + '_labels')

        # If not available, compute labels
        else:
            if verbose:
                print('Computing labels for ' + case_name)
            self.label_dict = compute_all_labels(self.flow_dict)
            path_to_mat = find_path_to_mat(case_name)
            save_dict_to_mat(path_to_mat[:-4] + '_labels', self.label_dict)

    def show_geometry(self):
        """
        Plot the geometry for the given case using a Line2D.
        :return: figure and axis object for the plot.
        """

        # Get points on the boundaries
        boundaries = get_boundaries(self)

        # Configure the plot
        xlim = [np.min(self.flow_dict['x']), np.max(self.flow_dict['x'])]
        ylim = [np.min(self.flow_dict['y']), np.max(self.flow_dict['y'])]
        linestyle = '-k'  # Solid black line

        # Plot boundaries
        fig, ax = plot.empty_plot()
        for boundary in boundaries:
            plot.lining(*boundary, linestyle=linestyle, xlim=xlim, ylim=ylim, append_to_fig_ax=(fig, ax))

        return fig, ax

    def show_flow(self, var_key, contour=True, contour_level=10, xlim=None, ylim=None, colorbar=False, show_geometry=True):
        """
        Display the requested variable in the physical coordinates.
        :param show_geometry: Print outline of the boundary.
        :param contour_level: Set amount of or specific levels.
        :param contour: Option for contour plot instead of scatter.
        :param ylim: List with shape [min_y, max_y]
        :param xlim: List with shape [min_x, max_x]
        :param var_key: Key in flow_dict for flow quantity.
        :return: void.
        """

        assert(var_key in self.flow_dict), "var_key is not a valid key for flow data: %r" % var_key

        title = var_key + ' in ' + self.case_name

        if self.flow_dict[var_key].shape[0] == 3:
            var = self.flow_dict[var_key][0]  # x-component of vector
        else:
            var = self.flow_dict[var_key]

        # 2D Geometry outline
        if show_geometry:
            fig, ax = self.show_geometry()
        else:
            fig, ax = plot.empty_plot()

        # Contour plot
        if contour:
            plot.contouring(self.flow_dict['x'].reshape(self.nx, self.ny),
                            self.flow_dict['y'].reshape(self.nx, self.ny),
                            var.reshape(self.nx, self.ny),
                            xlabel='x', ylabel='y',
                            title=title,
                            levels=contour_level,
                            colorbar=colorbar,
                            xlim=xlim, ylim=ylim,
                            append_to_fig_ax=(fig, ax))

        # Scatter plot
        else:
            plot.scattering(self.flow_dict['x'],
                            self.flow_dict['y'],
                            var,
                            scale=20, alpha=0.3,
                            colorbar=colorbar,
                            xlim=xlim, ylim=ylim,
                            xlabel='x', ylabel='y',
                            title=title,
                            append_to_fig_ax=(fig, ax))

    def show_label(self, label_key='non_negative', show_all=False, show_geometry=True, show_background=False, labelpos='center', only_positive=False, zoom_box=False):
        """
        Display the labels in the physical coordinates.
        :param show_geometry: Print outline of the boundary.
        :param label_key: Key for the label (see utilities.py)
        :param show_all: Display all labels.
        :return: void.
        """

        if show_all:
            for key in self.label_dict:

                # Create figure
                fig, ax = plot.empty_plot(figwidth=plot.latex_textwidth)
                boundaries, xlabel, ylabel, title, xlim, ylim, mask, seeding = plot.get_geometry_plot_data(self, return_seeding=True)

                plot.scattering(self.flow_dict['x'][~mask],
                                self.flow_dict['y'][~mask],
                                self.label_dict[key][~mask],
                                scale=10, alpha=1.0,
                                title=key,
                                cmap=plot.confusion_cmap,
                                append_to_fig_ax=(fig, ax))

                # 2D Geometry outline
                if show_geometry:
                    for boundary in boundaries:
                        plot.lining(*boundary,
                                    linestyle='-',
                                    color=plot.cblack,
                                    append_to_fig_ax=(fig, ax))
                    plot.set_labels_title(ax, xlabel, ylabel, title)
                    plot.set_limits(ax, xlim=xlim, ylim=ylim)
                else:
                    fig, ax = plot.empty_plot(figwidth=plot.latex_textwidth)

                # Background flow
                if show_background:
                    nx, ny = self.nx, self.ny
                    X, Y = self.flow_dict['x'].reshape(nx, ny), self.flow_dict['y'].reshape(nx, ny)
                    U, V = self.flow_dict['um'].reshape(nx, ny), self.flow_dict['vm'].reshape(nx, ny)
                    if "NACA" in self.case_name:
                        pass
                    else:
                        plot.streams(ax, X, Y, U, V, start_points=seeding, color=plot.cblack)
                    pass

        else:
            assert (label_key in self.label_dict), "label_key is not a valid key for labels: %r" % label_key

            # Create figure
            fig, ax = plot.empty_plot(figwidth=plot.latex_textwidth*0.9)
            boundaries, xlabel, ylabel, title, xlim, ylim, mask, seeding = plot.get_geometry_plot_data(self, return_seeding=True)

            # plot.scattering(self.flow_dict['x'][~mask],
            #                 self.flow_dict['y'][~mask],
            #                 self.label_dict[label_key][~mask],
            #                 scale=10,
            #                 alpha=1.0,
            #                 cmap=plot.confusion_cmap,
            #                 append_to_fig_ax=(fig, ax))
            X, Y, LABEL = self.flow_dict['x'][~mask], self.flow_dict['y'][~mask], self.label_dict[label_key][~mask]
            positive = LABEL == 1
            negative = LABEL == 0
            xs = [X[positive], X[negative]]
            ys = [Y[positive], Y[negative]]
            labels = [LABEL[positive], LABEL[negative]]
            zorders = [3, 2.9, 2.7, 2.8]
            if only_positive:
                xs = [xs[0]]
                ys = [ys[0]]
                labels = [labels[0]]
            for x, y, label, zorder in zip(xs, ys, labels, zorders):
                ax.scatter(x, y, s=10, c=label, cmap=plot.confusion_cmap, vmin=0, vmax=1, zorder=zorder)

            # Background flow
            if show_background:
                nx, ny = self.nx, self.ny
                X, Y = self.flow_dict['x'].reshape(nx, ny), self.flow_dict['y'].reshape(nx, ny)
                U, V = self.flow_dict['um'].reshape(nx, ny), self.flow_dict['vm'].reshape(nx, ny)
                if "NACA" in self.case_name:
                    pass
                else:
                    plot.streams(ax, X, Y, U, V, start_points=seeding, color=plot.cgrey)
                pass

            # 2D Geometry outline
            if show_geometry:
                for boundary in boundaries:
                    plot.lining(*boundary,
                                linestyle='-',
                                color=plot.cblack,
                                append_to_fig_ax=(fig, ax))
                plot.set_labels_title(ax, xlabel, ylabel, title)
                plot.set_limits(ax, xlim=xlim, ylim=ylim)
            else:
                fig, ax = plot.empty_plot(figwidth=plot.latex_textwidth)

            # Legend
            labels = [r"\textsf{True-Positive}", r"\textsf{True-Negative}", r"\textsf{False-Positive}", r"\textsf{False-Negative}"]
            # labels = ["Richtig-Positiv", "Richtig-Negativ", "Falsch-Positiv", "Falsch-Negativ"]
            label_values = [1, 0, 0.3, 0.7]

            cmap = plot.confusion_cmap
            clrs = cmap(label_values)

            legend_elements = list()
            for clr, label in zip(clrs[:2], labels[:2]):
                legend_elements.append(
                    plot.line2d([0], [0], marker='o', linestyle='', color=clr, markersize=4, markerfacecolor=clr, label=label))
            ax.legend(handles=legend_elements, loc=labelpos)

            # Zoom box
            if isinstance(zoom_box, list):
                for xi, yi in zip(zoom_box[0], zoom_box[1]):
                    # Create a Rectangle patch
                    rect = plot.patches.Rectangle((xi[0], yi[0]), abs(xi[1] - xi[0]), abs(yi[1] - yi[0]), linewidth=2.0,
                                                  edgecolor=plot.cblack, facecolor='none', zorder=10)
                    ax.add_patch(rect)

        return fig, ax

    def show_features(self, feature_key='tke', show_all=False, show_geometry=True):
        """
        Display the feature in the physical domain.
        :param show_geometry: Print outline of the boundary.
        :param feature_key: Key for the feature (see utilties.py and features.py)
        :param show_all: Display all features (WARNING >50 plots).
        :return: 1:success.
        """

        # Plot all features
        if show_all:
            for key in self.feature_dict:
                # 2D Geometry outline
                if show_geometry:
                    fig, ax = self.show_geometry()
                else:
                    fig, ax = plot.empty_plot()

                plot.scattering(self.flow_dict['x'],
                                self.flow_dict['y'],
                                self.feature_dict[key],
                                scale=5, alpha=1.0,
                                xlabel='x', ylabel='y', title=key,
                                colorbar=True,
                                append_to_fig_ax=(fig, ax))
        # Plot requested feature
        else:
            assert (feature_key in self.feature_dict), "feature_key is not a valid key for features: %r" % feature_key

            # 2D Geometry outline
            if show_geometry:
                fig, ax = self.show_geometry()
            else:
                fig, ax = plot.empty_plot()

            plot.scattering(self.flow_dict['x'],
                            self.flow_dict['y'],
                            self.feature_dict[feature_key],
                            scale=5, alpha=1.0,
                            xlabel='x/h', ylabel='y/h', title=feature_key,
                            cmap='viridis',
                            colorbar=True,
                            append_to_fig_ax=(fig, ax))

        return 1

    def show_profile(self, var_key, coord_key, loc_list):
        """
        Display profile along coordinate for any variable
        :param coord_key: Varying coordinate x or y.
        :param loc_list: Location of fixed coordinate.
        :param var_key: Variable to display on the profile.
        :return: void.
        """

        # Horizontal profile
        if coord_key == 'x':
            xlabel = coord_key
            ylabel = var_key

        # Vertical profile
        elif coord_key == 'y':
            xlabel = var_key
            ylabel = coord_key

        else:
            xlabel = 'error'
            ylabel = 'error'
            assert (coord_key == 'x' or coord_key == 'y'), 'Invalid key for coordinates, ' \
                                                           'must be x or y instead: %r' % coord_key

        # Get data for requested profile
        profile_data = get_profile_data(self, var_key, coord_key, loc_list)

        # Plot the profile
        plot.lining(*profile_data, xlabel=xlabel, ylabel=ylabel, title=self.case_name,
                    line_label=var_key + ' at y = ' + str(loc_list))
        return

    def show_integral_quantity(self, var_key, coord_key, xlog=False, ylog=False, xlim=None, ylim=None):
        """
        Plot an integral quantity along the given coordinate.
        :param ylim: List with shape [min_y, max_y]
        :param xlim: List with shape [min_x, max_x]
        :param ylog: Option for log-scaled y-axis.
        :param xlog: Option for log-scaled x-axis.
        :param var_key: Key for an integral quantity.
        :param coord_key: Key for any coordinate.
        :return: 1: success.
        """

        # For convenience
        coord = self.flow_dict[coord_key]
        var = self.flow_dict[var_key]
        nx = self.flow_dict['nx']
        ny = self.flow_dict['ny']

        if coord.shape == var.shape:
            plot.lining(coord, var,
                        xlog=xlog, ylog=ylog,
                        xlim=xlim, ylim=ylim,
                        xlabel=coord_key, ylabel=var_key,
                        line_label=None)
            return 1

        elif coord_key == 'x':
            plot.lining(coord[::ny], var,
                        xlog=xlog, ylog=ylog,
                        xlim=xlim, ylim=ylim,
                        xlabel=coord_key, ylabel=var_key,
                        line_label=None)
            return 1

        elif coord_key == 'y':
            plot.lining(coord[:ny], var,
                        xlog=xlog, ylog=ylog,
                        xlim=xlim, ylim=ylim,
                        xlabel=coord_key, ylabel=var_key,
                        line_label=None)
            return 1

        else:
            assert False, 'Cannot print variable and coordinate: %r and %r' % (var_key, coord_key)

    def show_anisotropy(self, sname=None, lumley=False, loc_key=None, loc_value=None):
        """
        Plot an barycentric mapping of the anisotropy tensor into an
        equilateral triangle. Optionally plot the Lumley triangle.
        :param loc_value: Value of the location e.g. x-coordinate.
        :param loc_key: Key of the location variable.
        :param lumley: Option for lumley triangle.
        :param save: Option to save the file under given string.
        :return: void.
        """

        # Check input arguments for profile data
        if loc_key and loc_value is None:
            assert loc_value, 'Only coordinate key given, missing location value (loc_value).'
        elif not loc_key and loc_value or loc_value == 0:
            assert loc_key, 'Only location value given, missing coordinate key (loc_coord).'

        # If location is requested, get and plot profile data
        elif loc_key and loc_value or loc_value == 0:
            # Index in profile data
            if loc_key == 'x':
                loc_idx = 1
            elif loc_key == 'y':
                loc_idx = 0
            else:
                assert False, 'Cannot handle given loc_key.'

            if lumley:
                profile_xi = get_profile_data(self, 'bij_xi', loc_key, loc_value)
                profile_eta = get_profile_data(self, 'bij_eta', loc_key, loc_value)
                plot.lumley_triangle(profile_xi[loc_idx],
                                     profile_eta[loc_idx],
                                     sname=sname)
            else:
                profile_eig1 = get_profile_data(self, 'bij_eig1', loc_key, loc_value)
                profile_eig2 = get_profile_data(self, 'bij_eig2', loc_key, loc_value)
                profile_eig3 = get_profile_data(self, 'bij_eig3', loc_key, loc_value)
                plot.barycentric_map(profile_eig1[loc_idx],
                                     profile_eig2[loc_idx],
                                     profile_eig3[loc_idx],
                                     sname=sname)

        # Plot all data points
        else:
            if lumley:
                plot.lumley_triangle(self.flow_dict['bij_xi'],
                                     self.flow_dict['bij_eta'],
                                     sname=sname)
            else:
                plot.barycentric_map(self.flow_dict['bij_eig1'],
                                     self.flow_dict['bij_eig2'],
                                     self.flow_dict['bij_eig3'],
                                     sname=sname)

    # def show_budgets(self, budget_key, coord_key, xlim=None, ylim=None):
    #     """
    #     Plot all budget terms for the tke in BL flow, tke or any tauij
    #     element along any coordinate.
    #     :param budget_key: Key for the budgets of interest e.g. k, uu, ..
    #     :param coord_key: Key for the coordinate.
    #     :param xlim: Limits of plot in x.
    #     :param ylim: Limits of plot in y.
    #     :return: 1: success, void: error.
    #     """
    #
    #     fig, ax = plot.empty_plot()
    #
    #     # Keys for k and tauij
    #     tau_k_keys = ['k', 'uu', 'vv', 'ww', 'uv', 'uw', 'vw']
    #
    #     # Find all budget keys for chosen budget
    #     for i, tau_k_key in enumerate(tau_k_keys):
    #         if tau_k_key == 'k' and self.case_name[:-5] == 'Bobke_FP_APG_LES':
    #             budget_keys = ALL_BUDGET_KEYS[i]
    #             break
    #         elif budget_key == 'k':
    #             budget_keys = ALL_BUDGET_KEYS[i+1]
    #             break
    #         # All tauij entries (uu, vv, ..)
    #         else:
    #             budget_keys = ALL_BUDGET_KEYS[i+1]
    #             break
    #
    #     # Plot budgets along coord_key
    #     for key in budget_keys:
    #         if key == 'diss_k':
    #             key = 'diss_rt'
    #         plot.lining(self.flow_dict[key], self.flow_dict[coord_key],
    #                     xlim=xlim, ylim=ylim,
    #                     ylabel=budget_key + ' budgets', xlabel=coord_key,
    #                     line_label=key,
    #                     append_to_fig_ax=(fig, ax))
    #
    #     return 1

#####################################################################
### Functions
#####################################################################
def compare_profiles(var_key, coord_key, loc, *cases, var_scale=1, xlim=None, ylim=None, xlog=False, ylog=False, append_to_fig_ax=(False, False)):
    """
    Compare a profile-plot setup between different flowCases.
    :param append_to_fig_ax: Append plot to another figure and axes.
    :param var_scale: Scaling factor for the variable.
    :param ylog: Optional logarithmic scale for y-axis.
    :param xlog: Optional logarithmic scale for x-axis.
    :param ylim: List with shape [min_y, max_y]
    :param xlim: List with shape [min_x, max_x]
    :param var_key: Any key within flowCase.flow_dict.
    :param coord_key: Varying coordinate.
    :param loc: Fixed location.
    :param cases: All flowCase objects.
    :return: Tuple of figure and axes.
    """

    # Generate new plot or append to given
    fig, ax = plot.check_append_to_fig_ax(append_to_fig_ax)

    for case in cases:
        profile_data = get_profile_data(case, var_key, coord_key, loc)

        # Horizontal profile
        if coord_key == 'x' or coord_key == 'y+':
            xlabel = coord_key
            ylabel = var_key
            scale_x = 1
            scale_y = var_scale

            # Vertical profile
        elif coord_key == 'y':
            xlabel = var_key
            ylabel = coord_key
            scale_x = var_scale
            scale_y = 1

        else:
            xlabel = 'error'
            ylabel = 'error'
            scale_x = 1
            scale_y = 1
            assert (coord_key == 'x' or coord_key == 'y'), 'Invalid key for coordinates, ' \
                                                           'must be x or y instead: %r' % coord_key

        plot.lining(*profile_data,
                    append_to_fig_ax=(fig, ax),
                    xlim=xlim, ylim=ylim,
                    xlog=xlog, ylog=ylog,
                    scale_x=scale_x, scale_y=scale_y,
                    xlabel=xlabel, ylabel=ylabel,
                    line_label=case.case_name)

    return fig, ax

def compare_integral_quantity(var_key, coord_key, *cases, xlog=False, ylog=False, xlim=None, ylim=None):
    """
    Compare an integral quantity along the given coordinate.
    :param ylim: List with shape [min_y, max_y]
    :param xlim: List with shape [min_x, max_x]
    :param ylog: Option for log-scaled y-axis.
    :param xlog: Option for log-scaled x-axis.
    :param var_key: Key for an integral quantity.
    :param coord_key: Key for any coordinate.
    :return: 1: success.
    """

    # Generate figure and subplot
    fig, ax = plot.empty_plot()

    for case in cases:
        # For convenience
        coord = case.flow_dict[coord_key]
        var = case.flow_dict[var_key]
        nx = case.flow_dict['nx']
        ny = case.flow_dict['ny']

        if coord.shape == var.shape:
            plot.lining(coord, var,
                        xlog=xlog, ylog=ylog,
                        xlim=xlim, ylim=ylim,
                        xlabel=coord_key, ylabel=var_key,
                        line_label=case.case_name,
                        append_to_fig_ax=(fig, ax))

        elif coord_key == 'x':
            plot.lining(coord[::ny], var,
                        xlog=xlog, ylog=ylog,
                        xlim=xlim, ylim=ylim,
                        xlabel=coord_key, ylabel=var_key,
                        line_label=case.case_name,
                        append_to_fig_ax=(fig, ax))

        elif coord_key == 'y':
            plot.lining(coord[:ny], var,
                        xlog=xlog, ylog=ylog,
                        xlim=xlim, ylim=ylim,
                        xlabel=coord_key, ylabel=var_key,
                        line_label=case.case_name,
                        append_to_fig_ax=(fig, ax))

        else:
            assert False, 'Cannot print variable and coordinate: %r and %r' % (var_key, coord_key)

    return 1

def compare_flow(var_key, *cases):
    """
    Compare distinct flow cases for a given variable on a
    grid of plots.
    :param var_key: Key of variable.
    :param cases: flowCase instances.
    :return: 1:success.
    """

    # Check valid key
    assert var_key in cases[0].flow_dict, "Invalid variable key given: %r" % var_key

    # Check required amount of subplots
    num_of_plots = len(cases)
    grid_length = int(num_of_plots ** 0.5) + 1

    # Create figure
    fig = plot.empty_plot(no_axis=True)

    # Add Title
    ax = fig.add_subplot(111)
    ax.set_title('Data basis with %r cases' % num_of_plots)
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)

    # Plot individual cases
    for i, case in enumerate(cases):
        ax = fig.add_subplot(grid_length, grid_length, i + 1)  # Add plot to grid
        ax.get_xaxis().set_visible(False)  # No axes
        ax.get_yaxis().set_visible(False)

        x = case.flow_dict['x']
        y = case.flow_dict['y']
        nx = case.flow_dict['nx']
        ny = case.flow_dict['ny']
        var = case.flow_dict[var_key]

        plot.contouring(x.reshape(nx, ny),
                        y.reshape(nx, ny),
                        var.reshape(nx, ny),
                        xlabel='', ylabel='',
                        colorbar=False,
                        append_to_fig_ax=(fig, ax))
    plot.show()

    return 1

# def test_init_case(flowcase_object):
#     """
#     Test initialisation of a flow case object.
#     :param flowcase_object: class flowcase_object type object.
#     :return: void.
#     """
#
#     # Check properties
#     assert(flowcase_object.num_of_points >= 1), "num_of_points is invalid: %r" % flowcase_object.num_of_points
#
#     return 1

# def test_flow_case(flowcase_object):
#     """
#     Run all tests on a flow case object.
#     :param flowcase_object: class flowcase_object type object.
#     :return: 1
#     """
#
#     print("BEGIN TEST flow case " + str(flowcase_object.case_name))
#
#     test_init_case(flowcase_object)
#     # test_features(flowcase_object)
#     # test_labels(flowcase_object)
#
#     print("COMPLETED TEST flow case " + str(flowcase_object.case_name))


import argparse
import numpy as np
import pickle
from keras.wrappers.scikit_learn import KerasClassifier
from keras.callbacks import ModelCheckpoint, TensorBoard
from keras.models import load_model
from sklearn.model_selection import cross_val_score
from sklearn.model_selection import KFold
from sklearn.metrics import accuracy_score
from sklearn.metrics import confusion_matrix
from phase_utils import print_cm
from phase_features_loader import PhaseFeaturesLoader
from phase_model_simple import model_simple
from phase_model_resnet import model_resnet
from phase_model_xgboost import model_xgboost
from imblearn.metrics import classification_report_imbalanced
from imblearn.combine import SMOTETomek, SMOTEENN

if __name__ == "__main__":
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("-a", "--action", choices=["train", "test"], default="train",
                        help="set the action, either training or test the dataset")
    parser.add_argument("--train_dataset", default="data/phase/ml_features_train.csv",
                        help="set the path to the training dataset")
    parser.add_argument("--test_dataset", default="data/phase/ml_features_test.csv",
                        help="set the path to the test dataset")
    parser.add_argument("-m", "--model", default=None,
                        help="set the path to the pre-trained model/weights")
    parser.add_argument("--cv", type=bool, default=False,
                        help="enable / disable a full cross validation with n_splits=10")
    parser.add_argument("-b", "--batch_size", type=int, default=256,
                        help="set the batch size)")
    parser.add_argument("-e", "--epochs", type=int, default=2000,
                        help="set the epochs number)")
    parser.add_argument("-l", "--layers", default="128 128 64 48 48 32 32 48 32 16",
                        help="set the hidden layers)")
    parser.add_argument("-d", "--dropout", type=float, default=0.1,
                        help="set the dropout)")
    parser.add_argument("-s", "--stations", default="URZ",
                        help="set the station name, it supports currently only LPAZ and URZ")
    parser.add_argument("-v", "--verbose", type=int, default=0,
                        help="set the verbosity)")
    parser.add_argument("-p", "--phase_length", default="URZ 6840 6840 6840 20520",
                        help="set the number of entries of phases per stations to be read from the dataset.\n" +
                             "The default is for the training, for the test use 'URZ 2280 2280 2280 6840, " +
                             "LPAZ 160 160 160 480'")

    args = parser.parse_args()

    # fix random seed for reproducibility
    seed = 7
    np.random.seed(seed)

    epochs = args.epochs
    train_dataset = args.train_dataset
    test_dataset = args.test_dataset
    phase_length = {}
    try:
        for p in args.phase_length.split(","):
            s = p.strip().split(" ")
            phase_length.update({s[0]:{"regP": int(s[1]), "regS": int(s[2]), "tele": int(s[3]), "N": int(s[4])}})
    except ValueError:
        print("It should be a list of a station name followed by four numbers.")
        exit(1)
    stations_lower = [station.lower() for station in sorted(phase_length.keys())]
    layers = []
    try:
        layers = [int(units) for units in args.layers.split(" ")]
    except ValueError:
        print("The layers should be a list of integer, delimited by a whitespace")
        exit(1)

    dropout = args.dropout
    batch_size = args.batch_size
    validation_split = 0.1
    if args.model is None:
        model_file_path = "results/phase_weights_best_s_{}_l_{}_d_{}.hdf5".\
            format("_".join(stations_lower), "_".join([str(layer) for layer in layers]), dropout)
    else:
        model_file_path = args.model

    model = model_xgboost

    if args.action == "train":
        # load train dataset
        pd = PhaseFeaturesLoader(filename=train_dataset, validation_split=validation_split,
                                 phase_length=phase_length, batch_size=batch_size)
        tensorboard = TensorBoard(log_dir='graph', histogram_freq=0, write_graph=True, write_images=True)
        checkpoint = ModelCheckpoint(model_file_path, monitor='acc', verbose=args.verbose,
                                     save_best_only=True, mode='max')
        if args.cv:
            train_x, train_y = pd.get_dataset()
            kfold = KFold(n_splits=10, shuffle=True, random_state=seed)
            estimator = KerasClassifier(build_fn=model, layers=layers, dropout=dropout,
                                        epochs=epochs, batch_size=500, verbose=args.verbose)
            results = cross_val_score(estimator, train_x, train_y, cv=kfold,
                                      fit_params={'callbacks':[checkpoint, tensorboard]})

            print("Baseline: %.2f%% (%.2f%%)" % (results.mean()*100, results.std()*100))
        else:
            model = model(layers=layers, dropout=dropout, layer_number=10)
            print(model)
            train_x, train_y = pd.get_dataset(expand_dim=False, y_onehot=False)
            sme = SMOTEENN(random_state=42)
            train_x_res, train_y_res = sme.fit_sample(train_x, train_y)
            #class_weight = {0:1, 1:1, 2:1, 3:1}
            # train_x = np.reshape(train_x, (train_x.shape[0], train_x[2]))
            # train_y = np.reshape(train_y, (train_y.shape[0], train_y[2]))
            model.fit(train_x_res, train_y_res)

            # save model to file
            pickle.dump(model, open("bst_grid-phase_smoteenn.pickle.dat", "wb"))

            print("Best accuracy obtained: {0}".format(model.best_score_))
            print("Parameters:")
            for key, value in model.best_params_.items():
                print("\t{}: {}".format(key, value))

            # load test dataset
            pd = PhaseFeaturesLoader(filename=test_dataset, phase_length=phase_length, batch_size=batch_size)
            test_x, test_y = pd.get_dataset(expand_dim=False, y_onehot=False)

            y_pred = model.predict(test_x)
            predictions = [np.round(value) for value in y_pred]
            # evaluate predictions
            accuracy = accuracy_score(test_y, predictions)
            print("Accuracy: %.2f%%" % (accuracy * 100.0))
    else:
        # load test dataset
        pd = PhaseFeaturesLoader(filename=test_dataset, phase_length=phase_length, batch_size=batch_size)
        test_x, test_y = pd.get_dataset(expand_dim=False, y_onehot=False)

        # load model & weight
        loaded_model = pickle.load(open("bst_grid-phase_smoteenn.pickle.dat", "rb"))
        print("Loaded model from disk with best params: {}".format(loaded_model.best_params_))

        y_pred = loaded_model.predict(test_x)

        predictions = [np.round(value) for value in y_pred]
        # evaluate predictions
        accuracy = accuracy_score(test_y, predictions)
        print("Accuracy: %.2f%%" % (accuracy * 100.0))
        class_name = ["regP", "regS", "tele", "N"]
        print(classification_report_imbalanced(test_y, predictions, target_names=class_name))


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

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

from siamban.core.config import cfg
from siamban.models.iou_loss import linear_iou

def get_cls_loss(pred, label, select):  # 通过这个来选出正负样本， label:(17500,) pred:(17500,2), select:(242,)
    if len(select.size()) == 0 or \
            select.size() == torch.Size([0]):
        return 0
    pred = torch.index_select(pred, 0, select)  # (242,2)
    label = torch.index_select(label, 0, select)  # (242,)
    return F.nll_loss(pred, label)

def select_cross_entropy_loss(pred, label):  # label:(28,25,25), pred:(28,25,25,2)
    pred = pred.view(-1, 2)  # (17500,2)
    label = label.view(-1)  # (17500,) 它的作用就是用来得到 下面的 pos 和 neg 的
    pos = label.data.eq(1).nonzero().squeeze().cuda()  # 正样本位置的索引
    neg = label.data.eq(0).nonzero().squeeze().cuda()  # 负样本位置的索引
    loss_pos = get_cls_loss(pred, label, pos)  #
    loss_neg = get_cls_loss(pred, label, neg)
    return loss_pos * 0.5 + loss_neg * 0.5

def weight_l1_loss(pred_loc, label_loc, loss_weight):
    if cfg.BAN.BAN:
        diff = (pred_loc - label_loc).abs()
        diff = diff.sum(dim=1)
    else:
        diff = None
    loss = diff * loss_weight
    return loss.sum().div(pred_loc.size()[0])

def select_iou_loss(pred_loc, label_loc, label_cls):
    label_cls = label_cls.reshape(-1)  # 17500
    pos = label_cls.data.eq(1).nonzero().squeeze().cuda()

    # pred_loc = pred_loc.permute(0, 2, 3, 1).reshape(-1, 4)  # (17500,4)
    pred_loc = torch.index_select(pred_loc, 0, pos)  # (242,4)

    # label_loc = label_loc.permute(0, 2, 3, 1).reshape(-1, 4)  # (17500,4)
    label_loc = torch.index_select(label_loc, 0, pos)  # (242,4)

    return linear_iou(pred_loc, label_loc)

def focal_loss(
    probs,
    targets,
    alpha: float = -1,
    gamma: float = 2,
    reduction: str = "none",
):
    ce_loss = F.binary_cross_entropy(
        probs, targets, reduction="none"
    )
    p_t = probs * targets + (1 - probs) * (1 - targets)
    loss = ce_loss * ((1 - p_t) ** gamma)

    if alpha >= 0:
        alpha_t = alpha * targets + (1 - alpha) * (1 - targets)
        loss = alpha_t * loss

    if reduction == "mean":
        loss = loss.mean()
    elif reduction == "sum":
        loss = loss.sum()

    return loss


"""Example how to make a Cube analysis from a 2D background model.
"""

from __future__ import absolute_import, division, print_function, unicode_literals
import numpy as np
from numpy.testing import assert_allclose
from astropy.tests.helper import assert_quantity_allclose
from astropy.coordinates import SkyCoord, Angle
from astropy.units import Quantity
from gammapy.extern.pathlib import Path
from ...data import DataStore
from ...utils.testing import requires_dependency, requires_data, pytest
from ...image import SkyMask
from .. import StackedObsCubeMaker
from ...background import OffDataBackgroundMaker
from .. import SkyCube
from ...utils.energy import Energy

def make_empty_cube(image_size, energy, center, data_unit=None):
    """
    Make an empty `SkyCube` from a given `SkyMask` and an energy binning.

    Parameters
    ----------
    image_size:int, Total number of pixel of the 2D map
    energy: Tuple for the energy axis: (Emin,Emax,nbins)
    center: SkyCoord of the source
    data_unit : str, Data unit.
    """
    def_image = dict()
    def_image["nxpix"] = image_size
    def_image["nypix"] = image_size
    def_image["binsz"] = 0.02
    def_image["xref"] = center.galactic.l.deg
    def_image["yref"] = center.galactic.b.deg
    def_image["proj"] = 'TAN'
    def_image["coordsys"] = 'GAL'
    def_image["unit"] = data_unit
    e_min, e_max, nbins = energy
    empty_cube = SkyCube.empty(emin=e_min.value, emax=e_max.value, enumbins=nbins, eunit=e_min.unit, mode='edges',
                               **def_image)
    return empty_cube

# Temp xfail for this: https://github.com/gammapy/gammapy/pull/899#issuecomment-281001655
@pytest.mark.xfail
@requires_dependency('reproject')
@requires_data('gammapy-extra')
def test_cube_pipe(tmpdir):
    tmpdir = str(tmpdir)
    outdir = tmpdir
    outdir2 = outdir + '/background'
    Path(outdir2).mkdir()

    ds = DataStore.from_dir("$GAMMAPY_EXTRA/datasets/hess-crab4-hd-hap-prod2")
    ds.copy_obs(ds.obs_table, tmpdir)
    data_store = DataStore.from_dir(tmpdir)

    # Create the background model from the 4 Crab observations
    bgmaker = OffDataBackgroundMaker(data_store, outdir=outdir2)

    bgmaker.select_observations(selection='all')
    bgmaker.group_observations()
    bgmaker.make_model("2D")
    bgmaker.save_models("2D")

    fn = outdir2 + '/group-def.fits'

    # New hdu table that contains the link to the background model
    hdu_index_table = bgmaker.make_total_index_table(
        data_store=data_store,
        modeltype='2D',
        out_dir_background_model=outdir2,
        filename_obs_group_table=fn
    )

    fn = outdir + '/hdu-index.fits.gz'
    hdu_index_table.write(fn, overwrite=True)

    center = SkyCoord(83.63, 22.01, unit='deg').galactic
    offset_band = Angle([0, 2.49], 'deg')

    ref_cube_images = make_empty_cube(image_size=250, energy=[Energy(0.5, "TeV"), Energy(100, "TeV"), 5], center=center)
    ref_cube_exposure = make_empty_cube(image_size=250, energy=[Energy(0.1, "TeV"), Energy(120, "TeV"), 80],
                                        center=center, data_unit="m2 s")
    ref_cube_skymask = make_empty_cube(image_size=250, energy=[Energy(0.5, "TeV"), Energy(100, "TeV"), 5],
                                       center=center)

    data_store = DataStore.from_dir(tmpdir)

    refheader = ref_cube_images.sky_image_ref.to_image_hdu().header
    exclusion_mask = SkyMask.read('$GAMMAPY_EXTRA/datasets/exclusion_masks/tevcat_exclusion.fits')
    exclusion_mask = exclusion_mask.reproject(reference=refheader)
    ref_cube_skymask.data = np.tile(exclusion_mask.data, (5, 1, 1))
    # Pb with the load psftable for one of the run that is not implemented yet...
    data_store.hdu_table.remove_row(14)

    # Cube Analysis
    cube_maker = StackedObsCubeMaker(empty_cube_images=ref_cube_images, empty_exposure_cube=ref_cube_exposure,
                                     offset_band=offset_band, data_store=data_store, obs_table=data_store.obs_table,
                                     exclusion_mask=ref_cube_skymask, save_bkg_scale=True)
    cube_maker.make_cubes(make_background_image=True, radius=10.)

    assert_allclose(cube_maker.counts_cube.data.sum(), 4898.0, atol=3)
    assert_allclose(cube_maker.bkg_cube.data.sum(), 4260.120595293951, atol=3)
    cube_maker.significance_cube.data[np.where(np.isinf(cube_maker.significance_cube.data))] = 0
    assert_allclose(np.nansum(cube_maker.significance_cube.data), 67613.24519908393, atol=3)
    assert_allclose(cube_maker.excess_cube.data.sum(), 637.8794047060486, atol=3)
    assert_quantity_allclose(np.nansum(cube_maker.exposure_cube.data), Quantity(4891844242766714.0, "m2 s"),
                             atol=Quantity(3, "m2 s"))
    assert_allclose(cube_maker.table_bkg_scale[0]["bkg_scale"], 0.8956177614218819)

    assert len(cube_maker.counts_cube.energies()) == 5
    assert len(cube_maker.bkg_cube.energies()) == 5
    assert len(cube_maker.significance_cube.energies()) == 5
    assert len(cube_maker.excess_cube.energies()) == 5
    assert len(cube_maker.exposure_cube.energies()) == 80


from typing import Union

import coreir
import hwtypes as ht

from magma.array import Array
from magma.bit import Bit
from magma.clock import Enable
from magma.common import ParamDict
from magma.bits import Bits, UInt, SInt
from magma.circuit import coreir_port_mapping, Circuit
from magma.conversions import as_bits, from_bits, bit
from magma.debug import magma_helper_function
from magma.interface import IO
from magma.generator import Generator
from magma.t import Type, Kind, In, Out, Direction
from magma.tuple import Tuple
from magma.clock import (
    AbstractReset, AsyncReset, AsyncResetN, Clock, get_reset_args
)
from magma.clock_io import ClockIO
from magma.primitives.mux import Mux
from magma.wire import wire

class _CoreIRRegister(Generator):
    """
    Internally used generator for CoreIR register primitive
    """
    def __init__(self, width, init=0, has_async_reset=False,
                 has_async_resetn=False):
        self.name = "reg_P"
        self.coreir_configargs = {"init": coreir.type.BitVector[width](init)}
        T = Bits[width]
        self.io = IO(I=In(T), CLK=In(Clock), O=Out(T))

        if has_async_resetn and has_async_reset:
            raise ValueError("Cannot have posedge and negedge asynchronous "
                             "reset")

        if has_async_reset:
            self.io += IO(arst=In(AsyncReset))
            self.name += "R"
            self.coreir_configargs["arst_posedge"] = True

        if has_async_resetn:
            self.io += IO(arst=In(AsyncResetN))
            self.name += "Rn"
            self.coreir_configargs["arst_posedge"] = False

        self.name += f"{width}"

        self.stateful = True
        self.primitive = True
        self.default_kwargs = {"init": coreir.type.BitVector[width](init)}
        self.coreir_genargs = {"width": width}
        self.renamed_ports = coreir_port_mapping

        self.coreir_name = "reg"
        if (has_async_reset or has_async_resetn):
            self.coreir_name += "_arst"

        self.verilog_name = "coreir_" + self.coreir_name
        self.coreir_lib = "coreir"

        def _simulate(self, value_store, state_store):
            cur_clock = value_store.get_value(self.CLK)

            if not state_store:
                state_store['prev_clock'] = cur_clock
                state_store['cur_val'] = ht.BitVector[width](init)

            if has_async_reset or has_async_resetn:
                cur_reset = value_store.get_value(self.arst)

            prev_clock = state_store['prev_clock']
            clock_edge = cur_clock and not prev_clock

            new_val = state_store['cur_val'].as_bool_list()

            if clock_edge:
                new_val = value_store.get_value(self.I)

            if has_async_reset and cur_reset or \
                    has_async_resetn and not cur_reset:
                new_val = ht.BitVector[width](init)

            state_store['prev_clock'] = cur_clock
            state_store['cur_val'] = ht.BitVector[width](new_val)
            value_store.set_value(self.O, new_val)
        self.simulate = _simulate

def _zero_init(T, init):
    if issubclass(T, Bits):
        return T(init)
    elif issubclass(T, Array):
        return T([_zero_init(T.T, init) for _ in range(len(T))])
    elif issubclass(T, Tuple):
        return T(*(_zero_init(t, init) for t in T.types()))
    return T(init)

def _get_T_from_init(init):
    if isinstance(init, (bool, ht.Bit)):
        return Bit
    if isinstance(init, int):
        return Bits[max(init.bit_length(), 1)]
    if isinstance(init, ht.UIntVector):
        return UInt[len(init)]
    if isinstance(init, ht.SIntVector):
        return SInt[len(init)]
    if isinstance(init, ht.BitVector):
        return Bits[len(init)]
    if isinstance(init, Type):
        return type(init)
    raise ValueError(f"Could not infer register type from {init}")

def _check_init_T(init, T):
    init_T = _get_T_from_init(init)
    if isinstance(init, int) and issubclass(T, Bits):
        # Allow int to be extended to width of T
        if len(init_T) > len(T):
            # Don't implicitly truncate
            return False
        return True
    if isinstance(init, int) and issubclass(T, Bit):
        # Allow int for bit
        if len(init_T) > 1:
            return False
        return True
    return init_T.is_wireable(T)

class AbstractRegister(Generator):
    """Abstract base class for all Register circuits. User-defined register
    types should implement this interface.
    """

    def get_enable(self, inst: Circuit) -> Enable:
        raise NotImplementedError()

class Register(AbstractRegister):
    def __init__(self, T: Kind = None,
                 init: Union[Type, int, ht.BitVector] = None,
                 reset_type: AbstractReset = None, has_enable: bool = False,
                 reset_priority: bool = True, name_map=ParamDict(I="I",
                                                                 O="O",
                                                                 CE="CE")):
        """
        T: The type of the value that is stored inside the register (e.g.
           Bits[5])

        init: (optional) A const value (i.e. init.const() == True) of type T or
              an int to be used as the initial value of the register.
              If no value is provided, the register will be initialized with 0

        has_enable: (optional) whether the register has an enable signal

        reset_type: (optional) The type of the reset port (also specifies the
                    semantic behavior of the reset signal)

        reset_priority: (optional) boolean flag choosing whether synchronous
                        reset (RESET or RESETN) has priority over enable

        name_map: (optional) ParamDict mapping default port names to new names
        """
        if T is None:
            if init is None:
                raise ValueError("User must provide type T or init value (from"
                                 " which T will be inferred)")
            T = _get_T_from_init(init)
        else:
            if not isinstance(T, Kind):
                raise TypeError(
                    f"Expected instance of Kind for argument T, not {type(T)}")
            if init is not None and not _check_init_T(init, T):
                raise ValueError(
                    f"Type {_get_T_from_init(init)} of init ({init}) does not "
                    f"match T ({T})"
                )

        (
            has_async_reset, has_async_resetn, has_reset, has_resetn
        ) = get_reset_args(reset_type)

        I_name = name_map.get("I", "I")
        O_name = name_map.get("O", "O")
        self.CE_name = CE_name = name_map.get("CE", "CE")
        self.io = IO(**{I_name: In(T), O_name: Out(T)})
        if has_enable:
            self.io += IO(**{CE_name: In(Enable)})
        self.io += ClockIO(has_async_reset=has_async_reset,
                           has_async_resetn=has_async_resetn,
                           has_reset=has_reset,
                           has_resetn=has_resetn)
        if init is None:
            init = 0

        if isinstance(init, int):
            init = _zero_init(T, init)

        self.init = init
        self.reset_type = reset_type

        coreir_init = init
        if isinstance(coreir_init, Type):
            coreir_init = as_bits(init)
        coreir_init = int(coreir_init)

        reg = _CoreIRRegister(T.flat_length(), init=coreir_init,
                              has_async_reset=has_async_reset,
                              has_async_resetn=has_async_resetn)()
        O = from_bits(T, reg.O)
        wire(O, getattr(self.io, O_name))

        I = getattr(self.io, I_name)
        if has_reset:
            reset_select = bit(self.io.RESET)
        elif has_resetn:
            reset_select = ~bit(self.io.RESETN)
        if has_enable:
            enable = getattr(self.io, CE_name)
        if (has_reset or has_resetn) and has_enable:
            if reset_priority:
                I = Mux(2, T)(name="enable_mux")(O, I, enable)
                I = Mux(2, T)()(I, init, reset_select)
            else:
                I = Mux(2, T)()(I, init, reset_select)
                I = Mux(2, T)(name="enable_mux")(O, I, enable)
        elif has_enable:
            I = Mux(2, T)(name="enable_mux")(O, I, enable)
        elif (has_reset or has_resetn):
            I = Mux(2, T)()(I, init, reset_select)
        reg.I @= as_bits(I)

    def get_enable(self, inst):
        return getattr(inst, self.CE_name, None)

@magma_helper_function
def register(value, **kwargs):
    T = type(value).qualify(Direction.Undirected)
    inst_kwargs = {}
    try:
        name = kwargs.pop("name")
        inst_kwargs["name"] = name
    except KeyError:
        pass
    ckt = Register(T, **kwargs)
    return ckt(**inst_kwargs)(value)


from project.car.muscle_car import MuscleCar
from project.car.sports_car import SportsCar
from project.driver import Driver
from project.race import Race

class Controller:
    VALID_CAR_TYPES = {
        "MuscleCar": MuscleCar,
        "SportsCar": SportsCar
    }

    def __init__(self):
        self.cars = []
        self.drivers = []
        self.races = []

    def create_car(self, car_type: str, model: str, speed_limit: int):
        if car_type in self.VALID_CAR_TYPES:
            if model in [car.model for car in self.cars]:
                raise Exception(f"Car {model} is already created!")
            self.cars.append(self.VALID_CAR_TYPES[car_type](model, speed_limit))
            return f"{car_type} {model} is created."

    def create_driver(self, driver_name: str):
        if driver_name in [driver.name for driver in self.drivers]:
            raise Exception(f"Driver {driver_name} is already created!")
        self.drivers.append(Driver(driver_name))
        return f"Driver {driver_name} is created."

    def create_race(self, race_name: str):
        if race_name in [r.name for r in self.races]:
            raise Exception(f"Race {race_name} is already created!")
        self.races.append(Race(race_name))
        return f"Race {race_name} is created."

    def add_car_to_driver(self, driver_name: str, car_type: str):
        try:
            driver = next(filter(lambda d: d.name == driver_name, self.drivers))
        except StopIteration:
            raise Exception(f"Driver {driver_name} could not be found!")

        try:
            car = list(filter(lambda c: c.__class__.__name__ == car_type and not c.is_taken, self.cars))[-1]
        except IndexError:
            raise Exception(f"Car {car_type} could not be found!")

        if driver.car:
            driver.car.is_taken = False
            old_car_model = driver.car.model
            car.is_taken = True
            driver.car = car
            return f"Driver {driver.name} changed his car from {old_car_model} to {car.model}."

        car.is_taken = True
        driver.car = car
        return f"Driver {driver_name} chose the car {car.model}."

    def add_driver_to_race(self, race_name: str, driver_name: str):
        try:
            race = next(filter(lambda r: r.name == race_name, self.races))
        except StopIteration:
            raise Exception(f"Race {race_name} could not be found!")

        try:
            driver = next(filter(lambda d: d.name == driver_name, self.drivers))
        except StopIteration:
            raise Exception(f"Driver {driver_name} could not be found!")

        if not driver.car:
            raise Exception(f"Driver {driver_name} could not participate in the race!")

        if driver in race.drivers:
            return f"Driver {driver_name} is already added in {race_name} race."

        race.drivers.append(driver)
        return f"Driver {driver_name} added in {race_name} race."

    def start_race(self, race_name: str):
        try:
            race = next(filter(lambda r: r.name == race_name, self.races))
        except StopIteration:
            raise Exception(f"Race {race_name} could not be found!")

        if len(race.drivers) < 3:
            raise Exception(f"Race {race_name} cannot start with less than 3 participants!")

        all_drivers = sorted(race.drivers, key=lambda x: x.car.speed_limit, reverse=True)
        final_result = []
        for i in range(3):
            current_driver = all_drivers[i]
            current_driver.number_of_wins += 1
            current_driver_result = f"Driver {current_driver.name} wins the {race_name} race" \
                                    f" with a speed of {current_driver.car.speed_limit}."
            final_result.append(current_driver_result)

        return '\n'.join(final_result)


import time

steps_program = 1

def EulerSolve(chessboard_size):
    x_move = [2, 1, -1, -2, -2, -1, 1, 2]
    y_move = [1, 2, 2, 1, -1, -2, -2, -1]

    print("Insert the X position of starting point from 0 - " + str(chessboard_size - 1) )
    starting_position_x = int(input())
    print("Insert the Y position of starting point from 0 - " + str(chessboard_size - 1) )
    starting_position_y = int(input())

    chessboard = [[-1 for i in range(chessboard_size)] for i in range(chessboard_size)]

    chessboard[starting_position_x][starting_position_y] = 0

    iteration = 1

    #checkHorseInChessboard(chessboard_size, position_x, position_y, x_move, y_move)
    #windstoffAlgorithm(starting_position_x,starting_position_y,chessboard,x_move, y_move,chessboard_size)
    start = time.time_ns()

    if (solveDPS(chessboard_size,chessboard,starting_position_x,starting_position_y,x_move,y_move, iteration)):
        printChessboard(chessboard,chessboard_size)
    else:
        print("NO SOLUTUTION FOR THIS INPUT")

    end = time.time_ns()
    total_time = end - start
    print("------------------------")
    print("TOTAL ATTEMPTS " + str(steps_program))
    print("TOTAL TIME: " + str(total_time / 1000000) + " milliseconds..")

def checkHorseInChessboard(chessboard_size, position_x, position_y, x_move, y_move,chessboard):
    layout = []
    for i in range(8):
        if moveValidation(position_x + x_move[i], position_y + y_move[i], chessboard_size,chessboard):
            layout.append([position_x + x_move[i], position_y + y_move[i]])
    return layout

def moveValidation(position_x, position_y, chessboard_size,chessboard):
    if position_x >= 0 and position_x < chessboard_size and position_y >= 0 and position_y < chessboard_size and chessboard[position_x][position_y] == -1:
        return True
    else:
        return False

def printChessboard(chessboard,chessboard_size):
    for x in range(chessboard_size):
        for y in range(chessboard_size):
            print(chessboard[x][y], end=" ")
        print("")

def bubblesort(list,x_move,y_move,chessboard):
    for iter_num in range(len(list)-1,0,-1):
      for idx in range(iter_num):
        first = checkHorseInChessboard(chessboard_size, list[idx][0], list[idx][1], x_move, y_move, chessboard)
        second = checkHorseInChessboard(chessboard_size, list[idx+1][0], list[idx+1][1], x_move, y_move, chessboard)
        if len(first)>len(second):
            temp = list[idx]
            list[idx] = list[idx+1]
            list[idx+1] = temp
            #print("CHANGE")

    return list

def solveDPS(chessboard_size,chessboard,position_x,position_y,x_move,y_move, iteration):
    if (iteration == chessboard_size ** 2):
        return True

    global steps_program

    positions = checkHorseInChessboard(chessboard_size, position_x, position_y, x_move, y_move, chessboard)
    positions = bubblesort(positions,x_move,y_move,chessboard)

    for i in positions:
        new_postion_x = i[0]
        new_postion_y = i[1]

        steps_program += 1
        if steps_program == 10000000:
            return False

        if (moveValidation(new_postion_x, new_postion_y, chessboard_size , chessboard)):
            chessboard[new_postion_x][new_postion_y] = iteration

            #FINAL CHECK
            if (solveDPS(chessboard_size,chessboard,new_postion_x,new_postion_y,x_move,y_move, iteration + 1)):
                return True

            chessboard[new_postion_x][new_postion_y] = -1

    #print("FALSE")
    return False

def finalCheck(chessboard,chessboard_size):

    checkmark = True
    for i in range(chessboard_size):
        for j in range(chessboard_size):
            if(chessboard[i][j] == 0):
                checkmark = False

    return checkmark

if __name__ == '__main__':
    print("Insert size of chessboard X: \n")
    chessboard_size = int(input())
    EulerSolve(chessboard_size)


from typing import Dict, Tuple
from unittest import TestCase

from google.protobuf.json_format import ParseDict

from cosmpy.common.utils import json_encode
from cosmpy.protos.cosmos.base.tendermint.v1beta1.query_pb2 import (
    GetBlockByHeightRequest,
    GetBlockByHeightResponse,
    GetLatestBlockRequest,
    GetLatestBlockResponse,
    GetLatestValidatorSetRequest,
    GetLatestValidatorSetResponse,
    GetNodeInfoRequest,
    GetNodeInfoResponse,
    GetSyncingRequest,
    GetSyncingResponse,
    GetValidatorSetByHeightRequest,
    GetValidatorSetByHeightResponse,
)
from cosmpy.tendermint.rest_client import CosmosBaseTendermintRestClient

from tests.helpers import MockRestClient

TYPE = {
    "@type": "type.googleapis.com/google.protobuf.Int32Value",
    "value": "42",
}

class CosmosBaseTendermintRestClientTestCase(TestCase):
    """Test case for CosmosBaseTendermintRestClient class."""

    REST_CLIENT = CosmosBaseTendermintRestClient

    def make_clients(
        self, response_content: Dict
    ) -> Tuple[MockRestClient, CosmosBaseTendermintRestClient]:
        """
        Make  mock client and rest client api for specific content.

        :param response_content: dict
        :return: rest client instance
        """
        mock_client = MockRestClient(json_encode(response_content).encode("utf-8"))
        rest_client = self.REST_CLIENT(mock_client)
        return mock_client, rest_client

    def test_GetNodeInfo(self):
        """Test GetNodeInfo method."""
        content = {
            "default_node_info": {
                "protocol_version": {
                    "p2p": "12",
                    "block": "12",
                    "app": "12",
                },
                "default_node_id": "string",
                "listen_addr": "string",
                "network": "string",
                "version": "string",
                "channels": "string",
                "moniker": "string",
                "other": {"tx_index": "string", "rpc_address": "string"},
            },
            "application_version": {
                "name": "string",
                "app_name": "string",
                "version": "string",
                "git_commit": "string",
                "build_tags": "string",
                "go_version": "string",
                "build_deps": [{"path": "string", "version": "string", "sum": "12"}],
            },
        }
        mock_client, rest_client = self.make_clients(content)
        expected_response = ParseDict(content, GetNodeInfoResponse())

        assert rest_client.GetNodeInfo(GetNodeInfoRequest()) == expected_response
        assert mock_client.last_base_url == "/cosmos/base/tendermint/v1beta1/node_info"

    def test_GetSyncing(self):
        """Test GetSyncing method."""
        content = {"syncing": True}
        mock_client, rest_client = self.make_clients(content)
        expected_response = ParseDict(content, GetSyncingResponse())

        assert rest_client.GetSyncing(GetSyncingRequest()) == expected_response
        assert mock_client.last_base_url == "/cosmos/base/tendermint/v1beta1/syncing"

    def test_GetLatestBlock(self):
        """Test GetLatestBlock method."""
        content = {
            "block_id": {
                "hash": "string",
                "part_set_header": {"total": 0, "hash": "string"},
            },
            "block": {
                "header": {
                    "version": {"block": "12", "app": "12"},
                    "chain_id": "string",
                    "height": "12",
                    "time": "2022-03-29T10:21:54.568Z",
                    "last_block_id": {
                        "hash": "string",
                        "part_set_header": {"total": 0, "hash": "string"},
                    },
                    "last_commit_hash": "string",
                    "data_hash": "string",
                    "validators_hash": "string",
                    "next_validators_hash": "string",
                    "consensus_hash": "string",
                    "app_hash": "string",
                    "last_results_hash": "string",
                    "evidence_hash": "string",
                    "proposer_address": "string",
                },
                "data": {"txs": ["string"]},
                "evidence": {"evidence": []},
                "last_commit": {
                    "height": "12",
                    "round": 0,
                    "block_id": {
                        "hash": "string",
                        "part_set_header": {"total": 0, "hash": "string"},
                    },
                    "signatures": [
                        {
                            "block_id_flag": "BLOCK_ID_FLAG_UNKNOWN",
                            "validator_address": "string",
                            "timestamp": "2022-03-29T10:21:54.569Z",
                            "signature": "string",
                        }
                    ],
                },
            },
        }
        mock_client, rest_client = self.make_clients(content)
        expected_response = ParseDict(content, GetLatestBlockResponse())

        assert rest_client.GetLatestBlock(GetLatestBlockRequest()) == expected_response
        assert (
            mock_client.last_base_url == "/cosmos/base/tendermint/v1beta1/blocks/latest"
        )

    def test_GetBlockByHeight(self):
        """Test GetBlockByHeight method."""
        content = {
            "block_id": {
                "hash": "string",
                "part_set_header": {"total": 0, "hash": "string"},
            },
            "block": {
                "header": {
                    "version": {"block": "12", "app": "12"},
                    "chain_id": "string",
                    "height": "12",
                    "time": "2022-03-29T10:27:01.686Z",
                    "last_block_id": {
                        "hash": "string",
                        "part_set_header": {"total": 0, "hash": "string"},
                    },
                    "last_commit_hash": "string",
                    "data_hash": "string",
                    "validators_hash": "string",
                    "next_validators_hash": "string",
                    "consensus_hash": "string",
                    "app_hash": "string",
                    "last_results_hash": "string",
                    "evidence_hash": "string",
                    "proposer_address": "string",
                },
                "data": {"txs": ["string"]},
                "evidence": {"evidence": []},
                "last_commit": {
                    "height": "12",
                    "round": 0,
                    "block_id": {
                        "hash": "string",
                        "part_set_header": {"total": 0, "hash": "string"},
                    },
                    "signatures": [
                        {
                            "block_id_flag": "BLOCK_ID_FLAG_UNKNOWN",
                            "validator_address": "string",
                            "timestamp": "2022-03-29T10:27:01.687Z",
                            "signature": "string",
                        }
                    ],
                },
            },
        }
        mock_client, rest_client = self.make_clients(content)
        expected_response = ParseDict(content, GetBlockByHeightResponse())

        assert (
            rest_client.GetBlockByHeight(GetBlockByHeightRequest(height=123))
            == expected_response
        )
        assert mock_client.last_base_url == "/cosmos/base/tendermint/v1beta1/blocks/123"

    def test_GetLatestValidatorSet(self):
        """Test GetLatestValidatorSet method."""
        content = {
            "block_height": "12",
            "validators": [
                {
                    "address": "string",
                    "pub_key": TYPE,
                    "voting_power": "12",
                    "proposer_priority": "12",
                }
            ],
            "pagination": {"next_key": "string", "total": "12"},
        }
        mock_client, rest_client = self.make_clients(content)
        expected_response = ParseDict(content, GetLatestValidatorSetResponse())

        assert (
            rest_client.GetLatestValidatorSet(GetLatestValidatorSetRequest())
            == expected_response
        )
        assert (
            mock_client.last_base_url
            == "/cosmos/base/tendermint/v1beta1/validatorsets/latest"
        )

    def test_GetValidatorSetByHeight(self):
        """Test GetValidatorSetByHeight method."""
        content = {
            "block_height": "12",
            "validators": [
                {
                    "address": "string",
                    "pub_key": TYPE,
                    "voting_power": "12",
                    "proposer_priority": "12",
                }
            ],
            "pagination": {"next_key": "string", "total": "12"},
        }
        mock_client, rest_client = self.make_clients(content)
        expected_response = ParseDict(content, GetValidatorSetByHeightResponse())

        assert (
            rest_client.GetValidatorSetByHeight(
                GetValidatorSetByHeightRequest(height=123)
            )
            == expected_response
        )
        assert (
            mock_client.last_base_url
            == "/cosmos/base/tendermint/v1beta1/validatorsets/123"
        )