""" Module for risk-exploiting landmark identification scenario
Environment will contain:
- Landmarks which can give reward and/or can cause an agent to fail with a probability based on distance (each has own function)
- A variable number of agents that can drop as agents are destroyed (at least 25 at start)
Agents will:
- know location location of landmarks but cannont directly observe the rewards/risk of landmarks except through indirect interaction
- Observe the location of failed/terminated agents
- be able to communicate with others within a certain distance (can communicate state, action, reward, and policy)
- have a high chance of failure on collision
"""

import numpy as np
from random import shuffle
from multiagent.scenario import BaseScenario
from particle_environments.mager.world import MortalAgent, HazardousWorld, RiskRewardLandmark
from particle_environments.mager.observation import format_observation
from particle_environments.common import is_collision, distance, delta_pos, delta_vel
from particle_environments.common import RadialPolynomialRewardFunction2D as RadialReward
from particle_environments.common import RadialBernoulliRiskFunction2D as RadialRisk
from rl_algorithms.scenariolearning import ScenarioHeuristicAgentTrainer

# Scenario Parameters
_MAX_COMMUNICATION_DISTANCE = np.inf
_AGENT_SIZE = 0.15
_LANDMARK_SIZE = 0.05
_AGENT_OBSERVATION_LEN = 5
_LANDMARK_OBSERVATION_LEN = 3
_NUM_AGENTS = 3

_LANDMARKS = []
_LANDMARKS.append(
    RiskRewardLandmark( risk_fn=RadialRisk(_LANDMARK_SIZE, 0.5), reward_fn=RadialReward(1.0, 0.0)))
_LANDMARKS.append(
    RiskRewardLandmark( risk_fn=None, reward_fn=RadialReward(1.0, 0.0)))
_LANDMARKS.append(
    RiskRewardLandmark( risk_fn=None, reward_fn=RadialReward(1.0, 0.0)))
_LANDMARKS.append(
    RiskRewardLandmark( risk_fn=None, reward_fn=RadialReward(1.0, 0.0)))
_N_LANDMARKS = len(_LANDMARKS)

class Scenario(BaseScenario):
    # static class
    num_agents = _NUM_AGENTS

    def make_world(self):
        world = HazardousWorld(collision_termination_probability=0.0)

        # observation-based communication
        world.dim_c = 0
        world.max_communication_distance = _MAX_COMMUNICATION_DISTANCE

        # collaborative rewards
        world.collaborative = True
        world.systemic_rewards = False
        world.identical_rewards = True

        # add landmarks
        world.landmarks = []
        for lm in _LANDMARKS:
            world.landmarks.append(lm)

        for i, landmark in enumerate(world.landmarks):
            landmark.name = 'landmark %d' % i
            landmark.collide = False
            landmark.movable = False
            landmark.size = _LANDMARK_SIZE
            # properties for landmarks
            if isinstance(landmark, RiskRewardLandmark) and landmark.is_hazard:
                #TODO: make colors heatmap of risk probability over all bounds
                landmark.color = np.array([landmark.risk_fn.get_failure_probability(0,0) + .1, 0, 0])
            else:
                landmark.color = np.array([0.25, 0.25, 0.25])

        # make initial conditions
        self.reset_world(world)
        return world

    def reset_world(self, world):
        # random properties for agents
        # add agents
        world.agents = [MortalAgent() for i in range(self.num_agents)]
        for i, agent in enumerate(world.agents):
            agent.name = 'agent %d' % i
            agent.collide = True
            agent.silent = True
            agent.terminated = False
            agent.size = _AGENT_SIZE
            agent.state.p_pos = np.random.uniform(-1, +1, world.dim_p)
            agent.state.p_vel = np.zeros(world.dim_p)
            agent.state.c = np.zeros(world.dim_c)
            agent.color = np.array([0.35, 0.35, 0.85])
            agent.previous_observation = None

        # shuffle landmarks to make sure hazard is not in same index
        shuffle(world.landmarks)
        for landmark in world.landmarks:
            landmark.state.p_pos = np.random.uniform(-1, +1, world.dim_p)
            landmark.state.p_vel = np.zeros(world.dim_p)
            landmark.hazard_tag = 0.0
            landmark.color = np.array([0.25, 0.25, 0.25])

    def benchmark_data(self, agent, world):
        collisions = 0
        occupied_landmarks = 0
        min_dists = 0
        for l in world.landmarks:
            dists = [np.linalg.norm(a.state.p_pos - l.state.p_pos) for a in world.agents]
            min_dists += min(dists)
            if min(dists) < 0.1:
                occupied_landmarks += 1
        if agent.collide:
            for a in world.agents:
                if is_collision(a, agent):
                    collisions += 1
        return (self.reward(agent, world), collisions, min_dists, occupied_landmarks)

    def done_callback(self, agent, world):
        ''' indicate a terminated agent as done '''
        # if agent.terminated:
        #     return True
        # else:
        #     return False
        return False

    def centralized_state_callback(self, world):
        ''' return the world truth to return state of all agents for training purposes
        Notes:
         - This CANNOT be used during execution, only training since it assumes the ability
           to centralize observations and map them to a state of the system
        '''
        state = []
        for ag in world.agents:
            state.extend([ag.terminated])
            state.extend(ag.state.p_pos)
            state.extend(ag.state.p_vel)

        for lm in world.landmarks:
            state.extend([lm.is_hazard])
            state.extend(lm.state.p_pos)
            state.extend(lm.state.p_vel)

        return state

    def reward(self, agent, world):
        # Agents are rewarded based on minimum agent distance to each landmark, penalized for collisions
        assert world.identical_rewards == True
        rew = 0
        for lm in [l for l in world.landmarks if not l.is_hazard]:
            # dists = [np.sqrt(np.sum(np.square(a.state.p_pos - l.state.p_pos))) for a in world.agents]
            dists = [distance(ag,lm) for ag in world.agents]
            rew -= min(dists)

        if agent.collide:
            for a in world.agents:
                if is_collision(a, agent):
                    rew -= 1
        return rew

    # def counterfactual_reward(self, central_state):
    #     ''' reward calculated based on a given centralized state vector, not the present world
    #     '''
    #     rew = 0.0
    #     for

    # def centralize_observations_to_partial_state_map(self, agent_observations):
    #     ''' use n-observations at a given time to generate a centralized, partial state
    #     Notes:
    #      - This CANNOT be used during execution, only training since it assumes a the ability
    #      to centralize observations and map them to a state of the system
    #     '''
    #     assert len(agent_observations) == _NUM_AGENTS
    #     state = []
    #     for ao in agent_observations:
    #         ao.extend(ao[0:5])

    def observation(self, agent, world):
        # get positions of all entities in this agent's reference frame

        def observe_agents(other_agent):
            ''' fill in information communicated/observed between agents
            '''
            # check if node is terminated
            is_terminated = 0
            if isinstance(other_agent, MortalAgent) and other_agent.terminated:
                is_terminated = 1

            dx = dy = dvx = dvy = 0.
            if not is_terminated:
                dx, dy = delta_pos(other_agent, agent)
                dvx, dvy = delta_vel(other_agent, agent)

            ag_obs = [is_terminated, dx, dy, dvx, dvy]
            assert(len(ag_obs) == _AGENT_OBSERVATION_LEN)
            return ag_obs

        def observe_landmarks(landmark):
            ''' fill in information observed about landmarks
            '''
            ld_obs = delta_pos(landmark, agent).tolist()

            # check if within observation range and is observable
            d = distance(landmark, agent)
            if d > world.max_communication_distance:
                ld_obs = [0.0]*len(ld_obs)

            # check if landmark is giving reward or hazard warning
            if d < landmark.size:
                if landmark.is_hazard:
                    landmark.hazard_tag = 1.0
                    landmark.color = np.array([1.1, 0, 0])
            ld_obs += [landmark.hazard_tag]

            assert(len(ld_obs) == _LANDMARK_OBSERVATION_LEN)
            return ld_obs

        landmark_positions = format_observation(observe = observe_landmarks,
                                                objects = world.landmarks,
                                                num_observations = len(world.landmarks),
                                                observation_size = _LANDMARK_OBSERVATION_LEN)

        agent_observations = format_observation(observe = observe_agents,
                                            objects = [a for a in world.agents if (a is not agent)],
                                            num_observations = self.num_agents-1,
                                            observation_size = _AGENT_OBSERVATION_LEN)

        new_obs = np.asarray([agent.terminated] + agent.state.p_vel.tolist() + agent.state.p_pos.tolist() + landmark_positions + agent_observations)

        # if agent.terminated:
        #     # if agent is terminated, return all zeros for observation
        #     # TODO: make this more efficient. Right now it does a lot of unnecessary calcs which are all
        #     #   then set to zero. Done this way to ensure consistant array size
        #     new_obs = 0.0*new_obs

        # append previous observation for velocity estimation
        # if agent.previous_observation is None:
        #     agent.previous_observation = 0.0*new_obs
        # obs = np.append(new_obs, agent.previous_observation)
        # agent.previous_observation = new_obs

        return new_obs

class ScenarioHeuristicComputer(ScenarioHeuristicAgentTrainer):
    ''' representation of an individual agent's embedded processor and memory tailor
    Notes:
     - This is meant to be used as a scenario-specific alternative to
     the more general purpose, scenario-agnostic "trainers". It can hold an agents model
     of the world (transition and reward functions), policy, and learning process, if any.
    '''
    def __init__(self, name, model, obs_shape_n, act_space_n, agent_index, args, **kwargs):
        ScenarioHeuristicAgentTrainer.__init__(self, name, model, obs_shape_n, act_space_n, agent_index, args)

        # record cumulative reward
        self.cumulative_reward = 0

        # behavior params
        bp_dist = self.get_initial_policy_distribution()
        bp = dict()
        for k in bp_dist:
            bp[k] = bp_dist[k][0]
        self.behavior_params = bp

        # observation range indices
        assert(_N_LANDMARKS==4)
        oir = dict()
        oir['terminated'] = [0,0]
        oir['velocity'] = [1,2]
        oir['position'] = [3,4]
        oir['landmarks'] = [5,4+_LANDMARK_OBSERVATION_LEN*_N_LANDMARKS]
        oir['agents'] = [oir['landmarks'][1]+1, oir['landmarks'][1] + _AGENT_OBSERVATION_LEN*(_NUM_AGENTS-1)]
        assert(oir['agents'][1] == obs_shape_n[0][0]-1)
        # oir['previous_observation'] = [oir['agents'][1]+1, 2*oir['agents'][1]+1]
        # assert(oir['previous_observation'][1] == obs_shape_n[0][0]-1)
        self.observation_index_ranges = oir

    def get_initial_policy_distribution(self):
        ''' method for "jumpstarting" monte carlo group distribution
        '''

        # (mean, std)
        bp_dist = dict()
        bp_dist['landmark_proximity_gain'] = (0.05, 0.1)
        bp_dist['agent_proximity_gain'] = (0.05, 0.1)

        return bp_dist

    def action(self, obs):
        ''' maps observation array to action forces in x,y directions
        Notes:
        '''

        # rename for ease of use
        oir = self.observation_index_ranges
        fx = 0.0
        fy = 0.0

        # find closest landmark
        targets = obs[oir['landmarks'][0]:oir['landmarks'][1]]
        target_index = oir['landmarks'][0]
        min_dist = np.linalg.norm(obs[target_index:target_index+1])
        for i in range(oir['landmarks'][0], oir['landmarks'][1]+1, _LANDMARK_OBSERVATION_LEN):
            cur_dist =np.linalg.norm(obs[i:i+1])
            if cur_dist < min_dist:
                min_dist = cur_dist
                target_index = i

        # force is proportional to distance to target
        dx_tar = obs[target_index]
        fx += self.behavior_params['landmark_proximity_gain']*dx_tar
        dy_tar = obs[target_index+1]
        fy += self.behavior_params['landmark_proximity_gain']*dy_tar

        # # force proportional to distance to agents
        # for i in range(oir['terminals'][1]+1, oir['terminals'][1]+1+21, 5):
        #     dx_agt = obs[i+1]
        #     dy_agt = obs[i+2]
        #     dsqr_agt_safe = max(_ZERO_THRESHOLD**2, dx_agt**2 + dy_agt**2)
        #     fx += -(dx_agt/dsqr_agt_safe)*self.behavior_params['agent_proximity_gain']
        #     fy += -(dy_agt/dsqr_agt_safe)*self.behavior_params['agent_proximity_gain']

        # check action force is valid
        assert(not np.isnan(fx))
        assert(not np.isnan(fy))

        # Due to (bizarre) formulation of action input in MultiAgentEnv __init__ and _set_action
        # we need to create a 5-element vector where element [1] and [3] are the x,y componets of
        # action force, respectivel
        act_force = np.zeros(5)
        act_force[1] = fx
        act_force[3] = fy

        return act_force

    def experience(self, obs, act, rew, new_obs, done, terminal):
        ''' Monte Carlo learning only record cumulative reward
        '''
        # record cumulative reward
        self.cumulative_reward += rew

    def preupdate(self):
        '''unused function handle compatibility with train.py
        '''
        pass

    def update(self, agents, t):
        '''unused function handle compatibility with train.py
        '''
        pass

    def group_policy_update(self, group_policy):
        '''update behavior parameters based on group policy
        '''
        for k in group_policy:
            self.behavior_params[k] = group_policy[k]['clambda']

        # reset cumulative reward
        self.cumulative_reward = 0


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

from model.utils.config import cfg
from model.faster_rcnn.faster_rcnn import _fasterRCNN
#from model.faster_rcnn.resnet import resnet

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import math
import torch.utils.model_zoo as model_zoo
import pdb
import numpy as np
import matplotlib.pyplot as plt

from model.rpn.bbox_transform import bbox_transform_inv, clip_boxes
from model.utils.net_utils import _affine_grid_gen

class GLCC(nn.Module):
    def __init__(self, FRCN, use_cuda, classes, net='res101', pretrained=False, class_agnostic=False):
        #resnet.__init__(self, classes=classes, num_layers=num_layers, pretrained=pretrained, class_agnostic=class_agnostic)
        super(GLCC, self).__init__()
        self.FRCN = FRCN
        if net == 'res101':
            self.fc_num = 2048
            self.roipool = 7
        elif net == 'res18':
            self.fc_num = 512
            self.roipool = 7
        elif net == 'alexnet':
            self.fc_num = 512
            self.roipool = 6
        elif net == 'vgg16':
            self.fc_num = 1024
            self.roipool = 7
        self.glcc_conv1 = nn.Conv2d(self.fc_num, 512, (1,1))
        self.glcc_fc1 = nn.Linear(self.roipool * self.roipool * 512, 1024)
        self.glcc_fc2 = nn.Linear(1024, 1024)
        self.glcc_fc_out = nn.Linear(1024, 4)
        self.use_cuda = use_cuda
        self.class_agnostic = class_agnostic
        self.classes = classes

    def forward(self, im_data, im_info, gt_boxes, num_boxes):
        batch_size = im_data.size(0)

        im_info = im_info.data
        gt_boxes = gt_boxes.data
        num_boxes = num_boxes.data

        rois, cls_prob, bbox_pred, rpn_loss_cls, rpn_loss_bbox, \
            RCNN_loss_cls, RCNN_loss_bbox, rois_label \
            = self.FRCN(im_data, im_info, gt_boxes, num_boxes)

        # get global and local region from Faster R-CNN

        base_feat = self.FRCN.RCNN_base(im_data)

        #print(rois.data.cpu().numpy())
        scores = cls_prob.data
        boxes = rois.data[:, :, 1:5]
        box_deltas = self.FRCN._bbox_pred.data

        if cfg.TRAIN.BBOX_NORMALIZE_TARGETS_PRECOMPUTED:
            if self.class_agnostic:
                if self.use_cuda > 0:
                    box_deltas = box_deltas.view(-1,4) * torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_STDS).cuda() + torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_MEANS).cuda()
                else:
                    box_deltas = box_deltas.view(-1,4) * torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_STDS) * torch.FlaotTensor(cfg.TRAIN.BBOX_NORMALIZE_MEANS)
                box_deltas = box_deltas.view(1, -1, 4)
            else:
                if self.use_cuda > 0:
                    box_deltas = box_deltas.view(-1,4) * torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_STDS).cuda() + torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_MEANS).cuda()
                else:
                    box_deltas = box_deltas.view(-1,4) * torhc.FlaotTensor(cfg.TRAIN.BBOX_NORMALIZE_STDS) + torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_MEANS)
                box_deltas = box_deltas.view(1, -1, 4 * len(self.classes))

        pred_boxes = bbox_transform_inv(boxes, box_deltas, 1)
        pred_boxes = clip_boxes(pred_boxes, im_info.data, 1)

        scores = scores.squeeze()
        pred_boxes = pred_boxes.squeeze()

        # get global region
        thresh = 0.00

        region_g = np.ndarray((0, 5))
        region_l = np.ndarray((0, 5))
        for j in range(1, 4):
            inds = torch.nonzero(scores[:,j]>=thresh).view(-1)
            inds_l = torch.nonzero(scores[:,j+3]>=thresh).view(-1)
            #print(inds)
            if inds.numel() > 0 and inds_l.numel() > 0:
                cls_scores = scores[:,j][inds]
                cls_scores_l = scores[:,j+3][inds_l]
                #print(cls_scores)
                #print(cls_scores_l)
                _, order = torch.sort(cls_scores, 0, True)
                _, order_l = torch.sort(cls_scores_l, 0, True)
                if self.class_agnostic:
                    cls_boxes = pred_boxes[inds]
                    cls_boxes_l = pred_boxes[inds_l]
                else:
                    cls_boxes = pred_boxes[inds][:,j*4:(j+1)*4]
                    cls_boxes_l = pred_boxes[inds_l][:,(j+3)*4:(j+4)*4]
                cls_dets = torch.cat((cls_boxes, cls_scores.unsqueeze(1)), 1)
                cls_dets_l = torch.cat((cls_boxes_l, cls_scores_l.unsqueeze(1)), 1)

                cls_dets = cls_dets[order]
                cls_dets_l = cls_dets_l[order_l]

                region_g = np.vstack((region_g, cls_dets[np.argmax(cls_dets[..., -1])]))
                region_l = np.vstack((region_l, cls_dets_l[np.argmax(cls_dets_l[..., -1])]))

        #print(cls_dets)
        #print(pred_boxes)

        # if true, then show detection global and local region
        if True:
            print(region_g)
            print(region_l)
            im = im_data.cpu().numpy()[0]
            im = np.transpose(im, (1,2,0))[..., ::-1]
            im -= im.min()
            im /= im.max()
            plt.imshow(im.astype(np.float))
            ax = plt.axes()
            ax.add_patch(
                plt.Rectangle((region_g[0, 0], region_g[0, 1]),
                              region_g[0, 2] - region_g[0, 0],
                              region_g[0, 3] - region_g[0, 1], fill=False,
                              edgecolor='red', linewidth=1))

            ax.add_patch(
                plt.Rectangle((region_l[0, 0], region_l[0, 1]),
                              region_l[0, 2] - region_l[0, 0],
                              region_l[0, 3] - region_l[0, 1], fill=False,
                              edgecolor='yellow', linewidth=1))
            plt.show()

        rois_g = np.zeros((1,1,5), dtype=np.float32)
        rois_g[0,0,1:5] = region_g[0, :4] / 16.
        rois_l = np.zeros((1,1,5), dtype=np.float32)
        rois_l[0,0,1:5] = region_l[0, :4] / 16.

        GPU = 0
        rois_g = torch.tensor(rois_g, dtype=torch.float).to(GPU)
        rois_l = torch.tensor(rois_l, dtype=torch.float).to(GPU)

        # global region
        if cfg.POOLING_MODE == 'crop':
            grid_xy = _affine_grid_gen(rois_g.view(-1, 5), base_feat.size()[2:], self.FRCN.grid_size)
            grid_yx = torch.stack([grid_xy.data[...,1], grid_xy.data[...,0]], 3).contiguous()
            pooled_feat_g = self.FRCN.RCNN_roi_crop(base_feat, Variable(grid_yx).detach())
            if cfg.CROP_RESIZE_WITH_MAX_POOL:
                pooled_feat_g = F.max_pool2d(pooled_feat_g, 2, 2)
        elif cfg.POOLING_MODE == 'align':
            pooled_feat_g = self.RCNN_roi_align(base_feat, rois_g.view(-1, 5))
        elif cfg.POOLING_MODE == 'pool':
            pooled_feat_g = self.RCNN_roi_pool(base_feat, rois_g.view(-1, 5))

        # local region
        if cfg.POOLING_MODE == 'crop':
            grid_xy = _affine_grid_gen(rois_l.view(-1, 5), base_feat.size()[2:], self.FRCN.grid_size)
            grid_yx = torch.stack([grid_xy.data[..., 1], grid_xy.data[...,0]], 3).contiguous()
            pooled_feat_l = self.FRCN.RCNN_roi_crop(base_feat, Variable(grid_yx).detach())
            if cfg.CROP_RESIZE_WITH_MAX_POOL:
                pooled_feat_l = F.max_pool2d(pooled_feat_l, 2, 2)
        elif cfg.POOLING_MODE == 'align':
            pooled_feat_l = self.RCNN_roi_align(base_feat, rois_l.view(-1,5))
        elif cfg.POOLING_MODE == 'pool':
            pooled_feat_l = self.RCNN_roi_pool(base_feat, rois_l.view(-1,5))

        #print(pooled_feat_g.cpu().detach().numpy().shape)
        x = torch.cat((pooled_feat_g, pooled_feat_l), dim=1)
        #print(x.cpu().detach().numpy().shape)
        x = self.glcc_conv1(x)
        x = F.relu(x)
        x = x.view(-1, self.roipool * self.roipool * 512)
        x = self.glcc_fc1(x)
        s = F.relu(x)
        x = nn.Dropout2d()(x)
        x = self.glcc_fc2(x)
        x = F.relu(x)
        x = nn.Dropout2d()(x)
        x = self.glcc_fc_out(x)

        return rois, cls_prob, bbox_pred, rpn_loss_cls, rpn_loss_bbox, RCNN_loss_cls, RCNN_loss_bbox, rois_label, x

    def train(self, mode=True):
        nn.Module.train(self, mode)
        if mode:
            self.FRCN.RCNN_base.eval()
            self.FRCN.RCNN_base[5].train()
            self.FRCN.RCNN_base[6].train()

            def set_bn_eval(m):
                classname = m.__class__.__name__
                if classname.find('BatchNorm') != -1:
                    m.eval()
            self.FRCN.RCNN_base.apply(set_bn_eval)
            self.FRCN.RCNN_top.apply(set_bn_eval)

    def _init_weights(self):
        def normal_init(m, mean, stddev, truncated=False):
            if truncated:
                m.weight.data.normal_().fmod_(2).mnul_(stddev).add_(mean)
            else:
                m.weight.data.normal_(mean, stddev)
                m.bias.data.zero_()

        normal_init(self.glcc_conv1, 0, 0.01, cfg.TRAIN.TRUNCATED)
        normal_init(self.glcc_fc1, 0, 0.01, cfg.TRAIN.TRUNCATED)
        normal_init(self.glcc_fc2, 0, 0.01, cfg.TRAIN.TRUNCATED)
        normal_init(self.glcc_fc_out, 0, 0.01, cfg.TRAIN.TRUNCATED)

    def create_architecture(self):
        self._init_modules()
        self._init_weights()


# -*- coding: utf-8 -*-

"""A module to perform optimal spectral extraction of SOSS time series observations"""

from functools import wraps
import os
from pkg_resources import resource_filename

from bokeh.plotting import show
from hotsoss import plotting as plt
from hotsoss import utils
from hotsoss import locate_trace as lt
import numpy as np

from . import decontaminate as dec
from . import summation as sm
from . import binning as bn
from . import sossfile as sf

def results_required(func):
    """A wrapper to check that the extraction has been run before a method can be executed"""
    @wraps(func)
    def _results_required(*args, **kwargs):
        """Check that the 'results' dictionary is not empty"""
        if not bool(args[0].results):
            print("No extraction found! Please run the 'extract' method first.")

        else:
            return func(*args, **kwargs)

    return _results_required

class SossExposure(object):
    """
    A class object to load a SOSS exposure file and extract and manipulate spectra
    """
    def __init__(self, filepath, name='My SOSS Observations', **kwargs):
        """
        Initialize the SOSS extraction object

        Parameters
        ----------
        filepath: str
            The path to the SOSS data, which must end with '_<ext>.fits',
            where 'ext' is ['uncal', 'ramp', 'rate', 'rateints', 'calints', 'x1dints']
        name: str
            The name of the exposure set
        """
        # Make sure the file exists
        if not os.path.exists(filepath):
            raise FileNotFoundError("{}: Invalid file".format(filepath))

        # Store attributes
        self.name = name
        self.time = None

        # Store empty SossFile objects
        self.uncal = sf.SossFile()  # 4D Uncalibrated raw input
        self.ramp = sf.SossFile() # 4D Corrected ramp data
        self.rate = sf.SossFile() # 2D Corrected countrate image
        self.rateints = sf.SossFile() # 3D Corrected countrate per integration
        self.calints = sf.SossFile() # 3D Calibrated data
        self.x1dints = sf.SossFile() # 1D Extracted spectra

        # Reset file levels
        self.levels = ['uncal', 'ramp', 'rate', 'rateints', 'calints', 'x1dints']

        # Load the file
        self.load_file(filepath)

        # Load the order throughput, wavelength calibration, and order mask files
        self.load_filters()
        self.order_masks = lt.order_masks(self.median)

        # Dictionary for the extracted spectra
        self.results = {}

        # Print uncal warning
        if self.uncal.file is not None:
            print("Looks like you have initialized an 'uncal' file! To pipeline process it, run 'SossExposure.uncal.calibrate()' method.")

    def calculate_order_masks(self):
        """
        Calculate the order masks from the median image
        """
        # Find the trace in all columns
        self.order_masks = lt.order_masks(self.median, save=True)

        print("New order masks calculated from median image.")

    def calibrate(self, ext='uncal', configdir=None, outdir=None, **kwargs):
        """
        Pipeline process a file in the exposure object

        Parameters
        ----------
        ext: str
            The extension to calibrate
        configdir: str
            The directory containing the configuration files
        outdir: str
            The directory to put the calibrated files into
        """
        if ext not in self.levels:
            raise ValueError("'{}' not valid extension. Please use {}".format(ext, self.levels))

        # Plot the appropriate file
        fileobj = getattr(self, ext)
        if fileobj.file is not None:
            new_files = fileobj.calibrate(configdir=configdir, outdir=outdir, **kwargs)

        # Load the calibrated files
        for level, file in new_files.items():
            self.load_file(file)

    @results_required
    def compare_results(self, idx=0, dtype='flux', names=None, order=None, draw=True):
        """
        Compare results of multiple extraction routines for a given frame

        Parameters
        ----------
        idx: int
            The frame index to plot
        dtype: str
            The data to plot, ['flux', 'counts']
        names: seqence (optional)
            The names of the extracted spectra to plot
        order: int (optional)
            The order to plot
        """
        # Get colors palette
        colors = utils.COLORS

        # Select the orders
        if isinstance(order, int):
            orders = [order]
        else:
            orders = [1, 2]

        if dtype == 'counts':
            ylabel = 'Counts [ADU/s]'
        else:
            ylabel = 'Flux Density [erg/s/cm2/A]'

        # Select the extractions
        fig = None
        if names is None:
            names = self.results.keys()

        for name in names:

            if name not in self.results:
                print("'{}' method not used for extraction. Skipping.".format(name))

            else:

                # Get the data dictionary and color
                result = self.results[name]
                color = next(colors)

                # Draw the figure
                data = result[dtype]
                wave = result['wavelength']
                flux = data[idx]
                fig = plt.plot_spectrum(wave, flux, fig=fig, legend=name, ylabel=ylabel, color=color, alpha=0.8)

        if fig is not None:
            if draw:
                show(fig)
            else:
                return fig

    @results_required
    def decontaminate(self, f277w_exposure):
        """
        Decontaminate the GR700XD+CLEAR orders with GR700XD+F277W order 1

        Parameters
        ----------
        f277w_exposure: SossExposure
            The F277W exposure object
        """
        # Check that the current filter is CLEAR
        if not self.filter == 'CLEAR':
            raise ValueError("filter = {}: Only a CLEAR exposure can be decontaminated".format(self.filter))

        # Check the dtype
        if not isinstance(f277w_exposure, type(self)):
            raise TypeError("{}: f277w_exposure must be of type {}".format(type(f277w_exposure), type(self)))

        # Check that the new filter is F277W
        new_filt = f277w_exposure.filter
        if new_filt != 'F277W':
            raise ValueError("filter = {}: Only an F277W exposure can be used for decontamination".format(new_filt))

        # Check that spectral extraction has been run on the F277W exposure
        if not bool(f277w_exposure.results):
            raise ValueError("Please run 'extract' method on F277W exposure before decontamination")

        # Run the decontamination
        self.results = dec.decontaminate(self.results, f277w_exposure.results)

    def extract(self, method="sum", ext='rateints', name=None, **kwargs):
        """
        Extract the 1D spectra from the time series data
        using the specified method

        Parameters
        ----------
        method: str
            The extraction method to use, ["bin", "sum"]
        ext: str
            The extension to extract
        name: str
            A name for the extraction results
        """
        # Validate the method
        valid_methods = ["bin", "sum"]
        if method not in valid_methods:
            raise ValueError("{}: Not a valid extraction method. Please use {}".format(method, valid_methods))

        # Set the extraction function
        func = bn.extract if method == "bin" else sm.extract

        # Get the requested data
        fileobj = getattr(self, ext)
        if fileobj.data is None:
            raise ValueError("No '{}' data to extract.".format(ext))

        # Run the extraction method, returning a dict with keys ['counts', 'wavelength', 'flux']
        result = func(fileobj.data, filt=self.filter, subarray=self.subarray, **kwargs)['final']
        result['method'] = method

        # Add the results to the table
        if name is None:
            name = method
        self.results[name] = result

    def _get_extension(self, ext, err=True):
        """
        Get the data for the given extension

        Parameters
        ----------
        ext: str
            The name of the extension
        err: bool
            Throw an error instead of printing
        """
        # Check the level is valid
        if ext not in self.levels:
            raise ValueError("'{}' not valid extension. Please use {}".format(ext, self.levels))

        # Get the data from the appropriate file
        fileobj = getattr(self, ext)
        if fileobj.file is None:
            loaded = [level for level in self.levels if getattr(self, level).file is not None]
            msg = "No data for '{0}' extension. Load `_{0}.fits' file or try {1}".format(ext, loaded)

            if err:
                raise ValueError(msg)
            else:
                print(msg)

        else:
            return fileobj

    @property
    def info(self):
        """Return some info about the observation"""
        # Pull out relevant attributes
        track = ['ncols', 'nrows', 'nints', 'ngrps', 'subarray', 'filter']
        settings = {key: val for key, val in self.__dict__.items() if key in track}

        # Get file info
        for level in self.levels:
            fileobj = getattr(self, level)
            file = None if fileobj is None else fileobj.file
            settings.update({level: file})

        return settings

    def load_file(self, filepath, **kwargs):
        """
        Load the data and headers from an observation file

        Parameters
        ----------
        filepath: str
            The path to the SOSS data
        """
        # Make sure the file exists
        if not os.path.exists(filepath):
            raise IOError("{} : Invalid file".format(filepath))

        # Determine processing level of input file
        if not any([filepath.endswith('_{}.fits'.format(level)) for level in self.levels]):
            raise ValueError('Not a recognized JWST file extension. Please use {}'.format(self.levels))

        # Save the filepath
        ext = filepath.split('_')[-1][:-5]
        fileobj = sf.SossFile(filepath)
        setattr(self, ext, fileobj)

        # Load the attributes
        self.nints = fileobj.nints
        self.ngrps = fileobj.ngrps
        self.nframes = fileobj.nframes
        self.nrows = fileobj.nrows
        self.ncols = fileobj.ncols
        self.filter = fileobj.filter
        self.frame_time = fileobj.frame_time
        self.subarray = fileobj.subarray
        self.wavecal = fileobj.wavecal
        self.median = fileobj.median
        self.time = fileobj.time

        print("'{}' file loaded from {}".format(ext, filepath))

    def load_filters(self):
        """
        Load the wavelength bins for orders 1, 2, and 3
        """
        self.filters = []

        # Pull out the throughput for the appropriate order
        for ord in [1, 2, 3]:
            file = resource_filename('hotsoss', 'files/GR700XD_{}.txt'.format(ord))
            if os.path.isfile(file):
                self.filters.append(np.genfromtxt(file, unpack=True))

    def plot_frames(self, ext='uncal', scale='linear', draw=True, **kwargs):
        """
        Plot a frame or all frames of any image data

        Parameters
        ----------
        ext: str
            The extension to plot
        idx: int
            The index of the frame to plot
        draw: bool
            Draw the figure instead of returning it
        """
        # Get the file object
        fileobj = self._get_extension(ext)

        # Make the plot
        fig = fileobj.plot(scale=scale)

        if draw:
            show(fig)
        else:
            return fig

    @results_required
    def plot_results(self, name=None, dtype='flux', time_fmt='mjd', draw=True):
        """
        Plot results of all integrations for the given extraction routine

        Parameters
        ----------
        name: str (optional)
            The name of the extracted data to plot
        dtype: str
            The data type to plot, ['flux', 'counts']
        time_fmt: str
            The astropy time format to use
        draw: bool
            Draw the figure instead of returning it
        """
        # Check dtype
        dtypes = ['flux', 'counts']
        if dtype not in dtypes:
            raise ValueError("{}: Please select dtype from {}".format(dtype, dtypes))

        # Check name
        names = list(self.results.keys())
        if name is not None and name not in names:
            raise ValueError("{}: Name not in results. Try {}".format(name, names))

        # Select the extractions
        fig = None
        if name is None:
            name = names[0]

        # Get the data
        result = self.results[name]
        data = result[dtype]
        wave = result['wavelength']
        time = self.time.to_value(time_fmt)

        # Draw the figure
        x = 'Wavelength [um]'
        y = 'Time [{}]'.format(time_fmt.upper())
        fig = plt.plot_time_series_spectra(data, wavelength=wave, time=time, ylabel=y, xlabel=x)

        if draw and fig is not None:
            show(fig)
        else:
            return fig

class SimExposure(SossExposure):
    """
    A test instance with the data preloaded
    """
    def __init__(self, subarray='SUBSTRIP256', filt='CLEAR', level='uncal', **kwargs):
        """
        Initialize the SOSS extraction object

        Parameters
        ----------
        subarray: str
            The desired subarray, ['SUBSTRIP96', 'SUBSTRIP256', 'FULL']
        filt: str
            The desired filter, ['CLEAR', 'F277W']
        level: str
            The desired level of pipeline processing, ['uncal', 'ramp']
        """
        # Get the file
        file = resource_filename('specialsoss', 'files/{}_{}_{}.fits'.format(subarray, filt, level))

        # Inherit from SossObs
        super().__init__(file, name='Simulated Observation', **kwargs)


import os
import sqlite3
import numpy as np
from astropy.time import Time
from rubin_sim.data import get_data_dir

__all__ = ["SeeingData"]

class SeeingData(object):
    """Read the seeing data from disk and return appropriate FWHM_500 value at a given time.
    This is for use in simulations only. Otherwise data would come from the EFD.

    Parameters
    ----------
    start_time : astropy.time.Time
        The time of the start of the simulation.
        The seeing database will be assumed to start on Jan 01 of the same year.
    seeing_db : str or None, optional
        The name of the seeing database.
        If None (default), this will use the simsee_pachon_58777_13.db file in the 'data' directory
        of this package.
        Other available seeing databases from sims_seeingModel include:
        seeing.db (the original, less-variable, 3 year seeing database)
        simsee_pachon_58777_13.db (the current default, 10 year, seeing database)
        simsee_pachon_58777_16.db (a similar, but slightly offset, 13 year seeing database)
        For more info on simsee_pachon_58777_*, see https://github.com/lsst/sims_seeingModel/issues/2
    offset_year : float, optional
        Offset into the cloud database by 'offset_year' years. Default 0.
    """
    def __init__(self, start_time, seeing_db=None, offset_year=0):
        self.seeing_db = seeing_db
        if self.seeing_db is None:
            self.seeing_db = os.path.join(get_data_dir(), 'site_models',
                                          'simsee_pachon_58777_13.db')

        # Seeing database starts in Jan 01 of the year of the start of the simulation
        year_start = start_time.datetime.year + offset_year
        self.start_time = Time('%d-01-01' % year_start, format='isot', scale='tai')

        self.seeing_dates = None
        self.seeing_values = None
        self.read_data()

    def __call__(self, time):
        """Get the FWHM_500 value for the specified time.

        Parameters
        ----------
        time : astropy.time.Time
            Time in the simulation for which to find the 'current' zenith seeing values.
            The difference between this time and the start_time, plus the offset,
            will be used to query the seeing database.

        Returns
        -------
        float
            The FWHM_500(") closest to the specified time.
        """
        delta_time = (time - self.start_time).sec
        # Find the date to look for in the time range of the data.
        # Note that data dates should not necessarily start at zero.
        dbdate = delta_time % self.time_range + self.min_time
        idx = np.searchsorted(self.seeing_dates, dbdate)
        # searchsorted ensures that left < date < right
        # but we need to know if date is closer to left or to right
        left = self.seeing_dates[idx - 1]
        right = self.seeing_dates[idx]
        if np.size(idx) == 1:
            if dbdate - left < right - dbdate:
                idx -= 1
        else:
            d1 = dbdate - left
            d2 = right - dbdate
            to_sub = np.where(d1 < d2)
            idx[to_sub] -= 1
        return self.seeing_values[idx]

    def read_data(self):
        """Read the seeing information from disk.

        The default behavior is to use the module stored database. However, an
        alternate database file can be provided. The alternate database file needs to have a
        table called *Seeing* with the following columns:

        seeingId
            int : A unique index for each seeing entry.
        s_date
            int : The time (in seconds) from the start of the simulation, for the seeing observation.
        seeing
            float : The FWHM of the atmospheric PSF (in arcseconds) at zenith.
        """
        with sqlite3.connect(self.seeing_db) as conn:
            cur = conn.cursor()
            query = "select s_date, seeing from Seeing order by s_date;"
            cur.execute(query)
            results = np.array(cur.fetchall())
            self.seeing_dates = np.hsplit(results, 2)[0].flatten()
            self.seeing_values = np.hsplit(results, 2)[1].flatten()
            cur.close()
        # Make sure seeing dates are ordered appropriately (monotonically increasing).
        ordidx = self.seeing_dates.argsort()
        self.seeing_dates = self.seeing_dates[ordidx]
        self.seeing_values = self.seeing_values[ordidx]
        self.min_time = self.seeing_dates[0]
        self.max_time = self.seeing_dates[-1]
        self.time_range = self.max_time - self.min_time

    def config_info(self):
        """Report information about configuration of this data.

        Returns
        -------
        OrderedDict
        """
        config_info = {}
        config_info['Start time for db'] = self.start_time
        config_info['Seeing database'] = self.seeing_db
        return config_info


import argparse
import logging
import numpy as np
import cv2
import torch
from os import makedirs
from os.path import isfile, isdir, join

from utils.log_helper import init_log, add_file_handler
from utils.bbox_helper import get_axis_aligned_bbox, cxy_wh_2_rect
from utils.load_helper import load_pretrain
from utils.benchmark_helper import load_dataset

from tools.test import siamese_init, siamese_track
from utils.config_helper import load_config
from utils.pyvotkit.region import vot_overlap, vot_float2str

def parse_range(arg):
    param = map(float, arg.split(','))
    return np.arange(*param)

def parse_range_int(arg):
    param = map(int, arg.split(','))
    return np.arange(*param)

parser = argparse.ArgumentParser(description='Finetune parameters for SiamMask tracker on VOT')
parser.add_argument('--arch', dest='arch', default='Custom', choices=['Custom', ],
                    help='architecture of pretrained model')
parser.add_argument('--resume', default='', type=str, required=True,
                    metavar='PATH',help='path to latest checkpoint (default: none)')
parser.add_argument('--config', dest='config',help='hyperparameter of SiamRPN in json format')
parser.add_argument('--mask', action='store_true', help='whether use mask output')
parser.add_argument('--refine', action='store_true', help='whether use mask refine output')
parser.add_argument('-v', '--visualization', dest='visualization', action='store_true',
                    help='whether visualize result')
parser.add_argument('--dataset', default='VOT2018', type=str,
                    metavar='DATASET', help='dataset')
parser.add_argument('-l', '--log', default="log_tune.txt", type=str,
                    help='log file')
parser.add_argument('--penalty-k', default='0.05,0.5,0.05', type=parse_range,
                    help='penalty_k range')
parser.add_argument('--lr', default='0.35,0.5,0.05', type=parse_range,
                    help='lr range')
parser.add_argument('--window-influence', default='0.1,0.8,0.05', type=parse_range,
                    help='window influence range')
parser.add_argument('--search-region', default='255,256,8', type=parse_range_int,
                    help='search region size')

args = parser.parse_args()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

def tune(param):
    regions = []  # result and states[1 init / 2 lost / 0 skip]
    # save result
    benchmark_result_path = join('result', param['dataset'])
    tracker_path = join(benchmark_result_path, (param['network_name'] +
                                                '_r{}'.format(param['hp']['instance_size']) +
                                                '_penalty_k_{:.3f}'.format(param['hp']['penalty_k']) +
                                                '_window_influence_{:.3f}'.format(param['hp']['window_influence']) +
                                                '_lr_{:.3f}'.format(param['hp']['lr'])).replace('.', '_'))  # no .
    if param['dataset'].startswith('VOT'):
        baseline_path = join(tracker_path, 'baseline')
        video_path = join(baseline_path, param['video'])
        result_path = join(video_path, param['video'] + '_001.txt')
    elif param['dataset'].startswith('OTB') or param['dataset'].startswith('DAVIS'):
        video_path = tracker_path
        result_path = join(video_path, param['video']+'.txt')

    if isfile(result_path):
        return

    try:
        if not isdir(video_path):
            makedirs(video_path)
    except OSError as err:
        print(err)

    with open(result_path, 'w') as f:  # Occupation
        f.write('Occ')

    global ims, gt, image_files
    if ims is None:
        print(param['video'] + '  Only load image once and if needed')
        ims = [cv2.imread(x) for x in image_files]
    start_frame, lost_times, toc = 0, 0, 0
    for f, im in enumerate(ims):
        tic = cv2.getTickCount()
        if f == start_frame:  # init
            cx, cy, w, h = get_axis_aligned_bbox(gt[f])
            target_pos = np.array([cx, cy])
            target_sz = np.array([w, h])
            state = siamese_init(im, target_pos, target_sz, param['network'], param['hp'], device=device)  # init tracker
            location = cxy_wh_2_rect(state['target_pos'], state['target_sz'])
            if param['dataset'].startswith('VOT'):
                regions.append(1)
            elif param['dataset'].startswith('OTB') or param['dataset'].startswith('DAVIS'):
                regions.append(gt[f])
        elif f > start_frame:  # tracking
            state = siamese_track(state, im, args.mask, args.refine, device=device)
            if args.mask:
                location = state['ploygon'].flatten()
            else:
                location = cxy_wh_2_rect(state['target_pos'], state['target_sz'])
            if param['dataset'].startswith('VOT'):
                if 'VOT' in args.dataset:
                    gt_polygon = ((gt[f][0], gt[f][1]),
                                  (gt[f][2], gt[f][3]),
                                  (gt[f][4], gt[f][5]),
                                  (gt[f][6], gt[f][7]))
                    if args.mask:
                        pred_polygon = ((location[0], location[1]), (location[2], location[3]),
                                        (location[4], location[5]), (location[6], location[7]))
                    else:
                        pred_polygon = ((location[0], location[1]),
                                        (location[0] + location[2], location[1]),
                                        (location[0] + location[2], location[1] + location[3]),
                                        (location[0], location[1] + location[3]))
                    b_overlap = vot_overlap(gt_polygon, pred_polygon, (im.shape[1], im.shape[0]))
                else:
                    b_overlap = 1

                if b_overlap:  # continue to track
                    regions.append(location)
                else:  # lost
                    regions.append(2)
                    lost_times += 1
                    start_frame = f + 5  # skip 5 frames
            else:
                regions.append(location)
        else:  # skip
            regions.append(0)
        toc += cv2.getTickCount() - tic

        if args.visualization and f >= start_frame:  # visualization (skip lost frame)
            if f == 0: cv2.destroyAllWindows()
            if len(gt[f]) == 8:
                cv2.polylines(im, [np.array(gt[f], np.int).reshape((-1, 1, 2))], True, (0, 255, 0), 3)
            else:
                cv2.rectangle(im, (gt[f, 0], gt[f, 1]), (gt[f, 0] + gt[f, 2], gt[f, 1] + gt[f, 3]), (0, 255, 0), 3)
            if len(location) == 8:
                location = np.int0(location)
                cv2.polylines(im, [location.reshape((-1, 1, 2))], True, (0, 255, 255), 3)
            else:
                location = [int(l) for l in location]  # bad support for OPENCV
                cv2.rectangle(im, (location[0], location[1]),
                              (location[0] + location[2], location[1] + location[3]), (0, 255, 255), 3)
            cv2.putText(im, str(f), (40, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2)  # frame id
            cv2.putText(im, str(lost_times), (40, 80), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)  # lost time

            cv2.imshow(param['video'], im)
            cv2.waitKey(1)
    toc /= cv2.getTickFrequency()
    print('Video: {:12s} Time: {:2.1f}s Speed: {:3.1f}fps Lost: {:d}'.format(param['video'], toc, f / toc, lost_times))

    with open(result_path, 'w') as f:
        for x in regions:
            f.write('{:d}\n'.format(x)) if isinstance(x, int) else \
                    f.write(','.join([vot_float2str("%.4f", i) for i in x]) + '\n')

def main():
    init_log('global', logging.INFO)
    if args.log != "":
        add_file_handler('global', args.log, logging.INFO)

    params = {'penalty_k': args.penalty_k,
              'window_influence': args.window_influence,
              'lr': args.lr,
              'instance_size': args.search_region}

    num_search = len(params['penalty_k']) * len(params['window_influence']) * \
        len(params['lr']) * len(params['instance_size'])

    print(params)
    print(num_search)

    cfg = load_config(args)
    if args.arch == 'Custom':
        from custom import Custom
        model = Custom(anchors=cfg['anchors'])
    else:
        model = models.__dict__[args.arch](anchors=cfg['anchors'])

    if args.resume:
        assert isfile(args.resume), '{} is not a valid file'.format(args.resume)
        model = load_pretrain(model, args.resume)
    model.eval()
    model = model.to(device)

    default_hp = cfg.get('hp', {})

    p = dict()

    p['network'] = model
    p['network_name'] = args.arch+'_'+args.resume.split('/')[-1].split('.')[0]
    p['dataset'] = args.dataset

    global ims, gt, image_files

    dataset_info = load_dataset(args.dataset)
    videos = list(dataset_info.keys())
    np.random.shuffle(videos)

    for video in videos:
        print(video)
        if isfile('finish.flag'):
            return

        p['video'] = video
        ims = None
        image_files = dataset_info[video]['image_files']
        gt = dataset_info[video]['gt']

        np.random.shuffle(params['penalty_k'])
        np.random.shuffle(params['window_influence'])
        np.random.shuffle(params['lr'])
        for penalty_k in params['penalty_k']:
            for window_influence in params['window_influence']:
                for lr in params['lr']:
                    for instance_size in params['instance_size']:
                        p['hp'] = default_hp.copy()
                        p['hp'].update({'penalty_k':penalty_k,
                                'window_influence':window_influence,
                                'lr':lr,
                                'instance_size': instance_size,
                                })
                        tune(p)

if __name__ == '__main__':
    main()
    with open('finish.flag', 'w') as f:  # Occupation
        f.write('finish')


# coding=utf-8

import keras
from keras.models import Sequential, model_from_json
from keras.layers import Dense, Dropout, Conv2D, MaxPooling2D, Flatten
from keras import backend as K
from keras.callbacks import ModelCheckpoint

from sklearn.model_selection import train_test_split
import h5py
import cv2
from PIL import Image
import numpy as np
from os.path import join
import os
import os.path

batch_size = 64
epochs = 30
input_shape = (40, 40, 1)
eng_list = list('abcdefghijklmnopqrstuvwxyz0123456789')
code_classes = 36
test_ratio = 0.3

x_list = []
y_list = []

# 读取图片
image_path = ''
img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
_, imgt = cv2.threshold(img, 140, 255, cv2.THRESH_BINARY_INV)

x_data = imgt / 255
y_data = ''

x_list.append(x_data)
y_list.append(y_data)

x_train, x_test, y_train, y_test = train_test_split(np.asarry(x_list), np.asarray(y_list),
                                                    test_size=test_ratio,
                                                    random_state=42)

# 网络搭建
model = Sequential()
model.add(Conv2D(16, kernel_size=(3, 3),
                 activation='relu',
                 input_shape=input_shape))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Flatten())
model.add(Dense(64, activation='relu'))
model.add(Dense(code_classes, activation='softmax'))

model.compile(loss=keras.losses.categorical_crossentropy,
              optimizer=keras.optimizers.Adadelta(),
              metrics=['accuracy'])

print(model.summary())
# 载入网络权重
model.load_weights('xxxx.hdf5')

# 回调函数，每个epoch结束后如果loss是历史最低的，则保存网络参数
checkpointer = ModelCheckpoint(filepath='xxxx.hdf5',
                               verbose=1,
                               save_best_only=True)
# 开始训练
model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs,
          verbose=1, validation_data=(x_test, y_test), callbacks=[checkpointer])

# 网络评价
score = model.evaluate(x_test, y_test, verbose=0)
print('Test loss: ', score[0])
print('Test accuracy: ', score[1])

# 保存网络的结构到json文件
json_string = model.to_json()
open('xxxx.json', 'w').write(json_string)


import tensorflow as tf
from tensorflow.keras.layers import Conv2D, GRU, Dense, Flatten, Input, Dense, Concatenate
import numpy as np
import random
from collections import deque
from tensorflow.keras.models import Model
import torch
import os

def Actor_network(obs_shape, num_agents, action_shape):
    obs = Input(shape=obs_shape)
    agent_onhot = Input(shape=num_agents)
    old_action = Input(shape=action_shape)
    input_embeding = Dense(64, activation="tanh")(tf.concat((obs, agent_onhot, old_action), axis=-1))
    #h = GRU(units=16, activation="relu", return_sequences=False)(input_embeding)
    hidden = Dense(32, activation="tanh")(input_embeding)
    pi = Dense(action_shape, activation="softmax")(hidden)
    model = Model(inputs=[obs, agent_onhot, old_action], outputs=pi)
    return model

def Critic_network(obs_shape, action_shape, num_agents, state_shape):
    obs = Input(shape=obs_shape)
    state = Input(shape=state_shape)
    agent_onhot = Input(shape=num_agents)
    old_actions = Input(shape=action_shape*num_agents)
    current_actions_without_agent = Input(shape=action_shape*num_agents)
    input_embeding = Dense(128, activation="tanh")(tf.concat((current_actions_without_agent, state, obs, agent_onhot, old_actions), axis=-1))
    hidden = Dense(64, activation="tanh")(input_embeding)
    q = Dense(action_shape)(hidden)
    model = Model(inputs=[current_actions_without_agent, state, obs, agent_onhot, old_actions], outputs=q)
    return model

class COMA():
    def __init__(self, n_actions, n_agents, state_shape, obs_shape):
        self.n_actions = n_actions
        self.n_agents = n_agents
        self.state_shape = state_shape
        self.obs_shape = obs_shape
        self.eval_policy = Actor_network(obs_shape=obs_shape, num_agents=n_agents, action_shape=n_actions)
        # 得到当前agent的所有可执行动作对应的联合Q值，得到之后需要用该Q值和actor网络输出的概率计算advantage
        self.eval_critic = Critic_network(obs_shape=obs_shape, action_shape=n_actions, num_agents=n_agents, state_shape=state_shape)
        self.target_critic = Critic_network(obs_shape=obs_shape, action_shape=n_actions, num_agents=n_agents, state_shape=state_shape)
        self.target_critic.trainable = False
        self.update_target_weights()
        print(self.eval_policy.summary())
        print(self.eval_critic.summary())
        print(self.target_critic.summary())
        self.eval_policy_optimizer = tf.keras.optimizers.RMSprop(learning_rate=1e-4)
        #self.eval_critic_optimizer = tf.keras.optimizers.RMSprop(learning_rate=1e-3)
        self.eval_critic_optimizer = tf.keras.optimizers.RMSprop(learning_rate=1e-3)

        self.summary_writer = tf.summary.create_file_writer("./log/coma_rmsprop_optimizer")

    def update_target_weights(self, tau=0.05):
        weights = self.eval_critic.get_weights()
        target_weights = self.target_critic.get_weights()
        for i in range(len(target_weights)):  # set tau% of target model to be new weights
            target_weights[i] = weights[i] * tau + target_weights[i] * (1 - tau)
        self.target_critic.set_weights(target_weights)

    def choose_action(self, obs, last_action, agent_num, epsilon, evaluate=False):
        # 传入的agent_num是一个整数，代表第几个agent，现在要把他变成一个onehot向量
        agent_onehot = np.zeros(self.n_agents)
        agent_onehot[agent_num] = 1

        obs = tf.convert_to_tensor([obs], dtype=tf.float32)
        agent_onehot = tf.convert_to_tensor([agent_onehot], dtype=tf.float32)
        last_action = tf.convert_to_tensor([last_action], dtype=tf.float32)

        pi = self.eval_policy.predict([obs, agent_onehot, last_action])
        if evaluate:
            action = np.argmax(pi[0])
        else:
            if np.random.rand(1) >= epsilon:  # epslion greedy
                action = np.random.randint(0, self.n_actions)
            else:
                action = np.argmax(pi[0])
        return action

    def learn(self, batch, max_episode_len, train_step, epsilon=0.9):  # train_step表示是第几次学习，用来控制更新target_net网络的参数
        # bacth中的每一项(n_episodes, episode_len, n_agents, 具体维度)
        # 我们采用最简化设计，n_episodes = 1， 即只采一轮数据就进行训练
        u, r, terminated = batch['u'], batch['r'], batch['terminated']

        with tf.GradientTape(persistent=True) as tape:
            # 根据经验计算每个agent的Ｑ值,从而跟新Critic网络。然后计算各个动作执行的概率，从而计算advantage去更新Actor。
            q_values, critic_loss = self._train_critic(batch, max_episode_len, train_step)  # 训练critic网络，并且得到每个agent的所有动作的Ｑ值
            action_prob = self._get_action_prob(batch, max_episode_len, epsilon)  # 每个agent的所有动作的概率
            q_taken = tf.squeeze(tf.gather(q_values, axis=3, indices=u, batch_dims=3), axis=3)  # 每个agent的选择的动作对应的Ｑ值
            pi_taken = tf.squeeze(tf.gather(action_prob, axis=3, indices=u, batch_dims=3), axis=3)  # 每个agent的选择的动作对应的概率
            log_pi_taken = tf.math.log(pi_taken)
            # 计算advantage
            baseline = tf.stop_gradient(tf.squeeze(tf.reduce_sum(tf.multiply(q_values, action_prob), axis=3, keepdims=True), axis=3))
            advantage = tf.stop_gradient(q_taken - baseline)
            policy_loss = - tf.reduce_sum(tf.multiply(advantage, log_pi_taken))

        critic_grads = tape.gradient(critic_loss, self.eval_critic.trainable_variables)
        self.eval_critic_optimizer.apply_gradients(zip(critic_grads, self.eval_critic.trainable_variables))

        policy_grads = tape.gradient(policy_loss, self.eval_policy.trainable_variables)
        self.eval_policy_optimizer.apply_gradients(zip(policy_grads, self.eval_policy.trainable_variables))

        if train_step > 0 and train_step % 10 == 0:
            self.update_target_weights()

        with self.summary_writer.as_default():
            tf.summary.scalar('Loss/polcy_loss', policy_loss, step=train_step)
            tf.summary.scalar('Loss/critic_loss', critic_loss, step=train_step)
        self.summary_writer.flush()

    def _get_action_prob(self, batch, max_episode_len, epsilon):
        episode_num = batch['o'].shape[0]

        action_prob = []
        obs = batch['o']  # (n_episodes, episode_len, n_agents, n_actions)
        agent_onhot = batch['a_onehot']  # (n_episodes, episode_len, n_agents, n_actions)
        u_onehot = batch['u_onehot']  # (n_episodes, episode_len, n_agents, n_actions)
        old_u_onehot = u_onehot[:, :-1]
        padded_old_u_onehot = np.zeros((u_onehot.shape[0], 1, u_onehot.shape[2], u_onehot.shape[3]))  # (n_episodes, episode_len, n_agents, 具体维度)

        old_u_onehot = np.concatenate((padded_old_u_onehot, old_u_onehot), axis=1)
        obs = np.reshape(obs, (obs.shape[0]*obs.shape[1]*obs.shape[2], obs.shape[3]))
        agent_onhot = np.reshape(agent_onhot, (agent_onhot.shape[0] * agent_onhot.shape[1] * agent_onhot.shape[2], agent_onhot.shape[3]))
        old_u_onehot = np.reshape(old_u_onehot, (old_u_onehot.shape[0] * old_u_onehot.shape[1] * old_u_onehot.shape[2], old_u_onehot.shape[3]))
        #model = Model(inputs=[obs, agent_onhot, old_action], outputs=pi)
        obs = tf.convert_to_tensor(obs, dtype=tf.float32)
        agent_onhot = tf.convert_to_tensor(agent_onhot, dtype=tf.float32)
        old_u_onehot = tf.convert_to_tensor(old_u_onehot, dtype=tf.float32)
        action_prob = self.eval_policy([obs, agent_onhot, old_u_onehot])
        # 把该列表转化成(episode个数, max_episode_len， n_agents，n_actions)的数组

        action_prob = ((1 - epsilon) * action_prob + tf.ones_like(action_prob) * epsilon / self.n_actions)
        # 因为上面把不能执行的动作概率置为0，所以概率和不为1了，这里要重新正则化一下。执行过程中Categorical会自己正则化。

        action_prob = action_prob/tf.reduce_sum(action_prob, axis=-1, keepdims=True)
        action_prob = tf.reshape(action_prob, (batch['u_onehot'].shape[0], batch['u_onehot'].shape[1], batch['u_onehot'].shape[2], batch['u_onehot'].shape[3]))
        return action_prob

    def _train_critic(self, batch, max_episode_len, train_step):
        # bacth中的每一项(n_episodes, episode_len, n_agents, 具体维度)
        # 我们采用最简化设计，n_episodes = 1， 即只采一轮数据就进行训练
        u, r, terminated = batch['u'], batch['r'], batch['terminated']
        u_next = u[:, 1:]
        padded_u_next = np.zeros((u.shape[0], 1, u.shape[2], u.shape[3]), dtype=np.int)#(n_episodes, episode_len, n_agents, 具体维度)
        u_next = np.concatenate((u_next, padded_u_next), axis=1)

        # 得到每个agent对应的Q值，维度为(episode个数, max_episode_len， n_agents，n_actions)
        # q_next_target为下一个状态-动作对应的target网络输出的Q值，没有包括reward
        # model = Model(inputs=[current_actions_without_agent, state, obs, agent_onhot, old_actions], outputs=q)
        '''
                           o=observations.copy(),
                           s=states.copy(),
                           u=actions_.copy(),
                           r=rewards.copy(),
                           o_next=observations_next.copy(),
                           s_next=states_next.copy(),
                           u_onehot=actions_onehots.copy(),
                           terminated=dones.copy(),
                           a_onehot = agents_onehot.copy()
        '''
        #get eval_critic inputs
        u_onehot = batch['u_onehot'] #(n_episodes, episode_len, n_agents, n_actions)
        old_u_onehot = u_onehot[:, :-1]
        padded_old_u_onehot = np.zeros((u_onehot.shape[0], 1, u_onehot.shape[2], u_onehot.shape[3]))  # (n_episodes, episode_len, n_agents, 具体维度)
        old_u_onehot = np.concatenate((padded_old_u_onehot, old_u_onehot), axis=1)

        next_u_onehot = u_onehot[:, 1:]
        padded_next_u_onehot = np.zeros((u_onehot.shape[0], 1, u_onehot.shape[2], u_onehot.shape[3]))  # (n_episodes, episode_len, n_agents, 具体维度)
        next_u_onehot = np.concatenate((padded_next_u_onehot, next_u_onehot), axis=1)

        u_onehot = np.reshape(u_onehot, (batch['u_onehot'].shape[0]*batch['u_onehot'].shape[1],
                                     batch['u_onehot'].shape[2]*batch['u_onehot'].shape[3]))
        u_onehot = np.repeat(np.expand_dims(u_onehot, axis=1), repeats=self.n_agents, axis=1)
        #(n_episodes*episode_len, n_agents, n_agents*n_actions)
        u_onehot_for_eval_critic = u_onehot
        u_onehot_for_target_critic = u_onehot

        u_onehot_for_eval_critic_set_agent_o = np.ones(shape=(u_onehot_for_eval_critic.shape[0], u_onehot_for_eval_critic.shape[1], u_onehot_for_eval_critic.shape[2]))
        for i in range(self.n_agents):
            u_onehot_for_eval_critic_set_agent_o[:, i, i*self.n_actions:i*self.n_actions+self.n_actions] = 0
        u_onehot_for_eval_critic = u_onehot_for_eval_critic * u_onehot_for_eval_critic_set_agent_o
        #u_onehot_for_eval_critic = tf.convert_to_tensor(u_onehot_for_eval_critic, dtype=tf.float32)

        old_u_onehot = np.reshape(old_u_onehot, (batch['u_onehot'].shape[0] * batch['u_onehot'].shape[1],
                                             batch['u_onehot'].shape[2] * batch['u_onehot'].shape[3]))
        old_u_onehot = np.repeat(np.expand_dims(old_u_onehot, axis=1), repeats=self.n_agents, axis=1)
        # (n_episodes*episode_len, n_agents, n_agents*n_actions)

        current_actions_without_agent = np.reshape(u_onehot_for_eval_critic, (u_onehot_for_eval_critic.shape[0]*u_onehot_for_eval_critic.shape[1], u_onehot_for_eval_critic.shape[2]))
        state = np.reshape(batch['s'], (batch['s'].shape[0] * batch['s'].shape[1] * batch['s'].shape[2], -1))
        obs = np.reshape(batch['o'], (batch['o'].shape[0] * batch['o'].shape[1] * batch['o'].shape[2], -1))
        agent_onhot = np.reshape(batch['a_onehot'], (batch['a_onehot'].shape[0] * batch['a_onehot'].shape[1] * batch['a_onehot'].shape[2], -1))
        old_actions = np.reshape(old_u_onehot, (old_u_onehot.shape[0]*old_u_onehot.shape[1], old_u_onehot.shape[2]))

        # get target_critic inputs
        next_u_onehot = np.reshape(next_u_onehot, (batch['u_onehot'].shape[0] * batch['u_onehot'].shape[1],
                                     batch['u_onehot'].shape[2] * batch['u_onehot'].shape[3]))
        next_u_onehot = np.repeat(np.expand_dims(next_u_onehot, axis=1), repeats=self.n_agents, axis=1)
        # (n_episodes*episode_len, n_agents, n_agents*n_actions)
        next_u_onehot = next_u_onehot * u_onehot_for_eval_critic_set_agent_o

        target_current_actions_without_agent = np.reshape(next_u_onehot, (next_u_onehot.shape[0] * next_u_onehot.shape[1], next_u_onehot.shape[2]))
        target_state = np.reshape(batch['s_next'], (batch['s_next'].shape[0] * batch['s_next'].shape[1] * batch['s_next'].shape[2], -1))
        target_obs = np.reshape(batch['o_next'], (batch['o_next'].shape[0] * batch['o_next'].shape[1] * batch['o_next'].shape[2], -1))
        target_agent_onhot = np.reshape(batch['a_onehot'], (batch['a_onehot'].shape[0] * batch['a_onehot'].shape[1] * batch['a_onehot'].shape[2], -1))
        target_old_actions = np.reshape(u_onehot_for_target_critic, (u_onehot_for_target_critic.shape[0] * u_onehot_for_target_critic.shape[1], u_onehot_for_target_critic.shape[2]))

        current_actions_without_agent = tf.convert_to_tensor(current_actions_without_agent, dtype=tf.float32)
        state = tf.convert_to_tensor(state, dtype=tf.float32)
        obs = tf.convert_to_tensor(obs, dtype=tf.float32)
        agent_onhot = tf.convert_to_tensor(agent_onhot, dtype=tf.float32)
        old_actions = tf.convert_to_tensor(old_actions, dtype=tf.float32)
        target_current_actions_without_agent = tf.convert_to_tensor(target_current_actions_without_agent, dtype=tf.float32)
        target_state = tf.convert_to_tensor(target_state, dtype=tf.float32)
        target_obs = tf.convert_to_tensor(target_obs, tf.float32)
        target_agent_onhot = tf.convert_to_tensor(target_agent_onhot, dtype=tf.float32)
        target_old_actions = tf.convert_to_tensor(target_old_actions, dtype=tf.float32)

        # (n_episodes*episode_len*n_agents, n_actions)
        q_evals = self.eval_critic([current_actions_without_agent, state, obs, agent_onhot, old_actions])
        q_next_target = self.target_critic([target_current_actions_without_agent, target_state, target_obs, target_agent_onhot, target_old_actions])
        q_evals = tf.reshape(q_evals, (batch['o_next'].shape[0], batch['o_next'].shape[1], batch['o_next'].shape[2], self.n_actions))
        q_next_target = tf.reshape(q_next_target, (batch['o_next'].shape[0], batch['o_next'].shape[1], batch['o_next'].shape[2], self.n_actions))
        # (n_episodes, episode_len, n_agents, n_actions)

        q_values = q_evals  # 在函数的最后返回，用来计算advantage从而更新actor
        # 取每个agent动作对应的Q值，并且把最后不需要的一维去掉，因为最后一维只有一个值了

        # u: (n_episodes, episode_len, n_agents, 1)
        # q_evals: (n_episodes, episode_len, n_agents, n_actions)
        q_evals = tf.squeeze(tf.gather(params=q_evals, indices=u, axis=3, batch_dims=3), axis=3)# (n_episodes, episode_len, n_agents)

        q_next_target = tf.squeeze(tf.gather(params=q_next_target, indices=u_next, axis=3, batch_dims=3), axis=3)# (n_episodes, episode_len, n_agents)

        targets = self.td_lambda_target(batch, max_episode_len, q_next_target)#(episode_num, max_episode_len, n_agents)

        td_error = tf.stop_gradient(tf.convert_to_tensor(targets, dtype=tf.float32)) - q_evals#(episode_num, max_episode_len, n_agents)

        # 不能直接用mean，因为还有许多经验是没用的，所以要求和再比真实的经验数，才是真正的均值
        loss = tf.reduce_sum(tf.square(td_error))
        # print('Loss is ', loss)

        return q_values, loss

    def td_lambda_target(self, batch, max_episode_len, q_targets, gamma=0.99, td_lambda=0.8):  # 用来通过TD(lambda)计算y
        # batch维度为(episode个数, max_episode_len， n_agents，n_actions)
        # q_targets维度为(episode个数, max_episode_len， n_agents)
        episode_num = batch['o'].shape[0]
        #mask = (1 - batch["padded"].float()).repeat(1, 1, self.n_agents)  # 用来把那些填充的经验的TD-error置0，从而不让它们影响到学习
        terminated = (1 - batch["terminated"]).repeat(repeats= self.n_agents, axis=2)# 用来把episode最后一条经验中的q_target置0
        # 把reward维度从(episode个数, max_episode_len, 1)变成(episode个数, max_episode_len, n_agents)
        r = batch['r'].repeat(repeats= self.n_agents, axis=2)
        # 计算n_step_return

        '''
        1. 每条经验都有若干个n_step_return，所以给一个最大的max_episode_len维度用来装n_step_return
        最后一维,第n个数代表 n+1 step。
        2. 因为batch中各个episode的长度不一样，所以需要用mask将多出的n-step return置为0，
        否则的话会影响后面的lambda return。第t条经验的lambda return是和它后面的所有n-step return有关的，
        如果没有置0，在计算td-error后再置0是来不及的
        3. terminated用来将超出当前episode长度的q_targets和r置为0
        '''
        n_step_return = np.zeros((episode_num, max_episode_len, self.n_agents, max_episode_len))
        for transition_idx in range(max_episode_len - 1, -1, -1):
            # 最后计算1 step return
            n_step_return[:, transition_idx, :, 0] = (r[:, transition_idx] + gamma * q_targets[:, transition_idx] * terminated[:, transition_idx])
            # 经验transition_idx上的obs有max_episode_len - transition_idx个return, 分别计算每种step return
            # 同时要注意n step return对应的index为n-1
            for n in range(1, max_episode_len - transition_idx):
                # t时刻的n step return =r + gamma * (t + 1 时刻的 n-1 step return)
                # n=1除外, 1 step return =r + gamma * (t + 1 时刻的 Q)
                n_step_return[:, transition_idx, :, n] = (r[:, transition_idx] + gamma *
                                                          n_step_return[:, transition_idx + 1, :, n - 1])
            # 计算lambda return
        '''
        lambda_return 维度为(episode个数, max_episode_len， n_agents)，每条经验中，每个agent都有一个lambda return
        '''
        lambda_return = np.zeros((episode_num, max_episode_len, self.n_agents))
        for transition_idx in range(max_episode_len):
            returns = np.zeros((episode_num, self.n_agents))
            for n in range(1, max_episode_len - transition_idx):
                returns += pow(td_lambda, n - 1) * n_step_return[:, transition_idx, :, n - 1]
            lambda_return[:, transition_idx] = (1 - td_lambda) * returns + \
                                               pow(td_lambda, max_episode_len - transition_idx - 1) * \
                                               n_step_return[:, transition_idx, :, max_episode_len - transition_idx - 1]
        return lambda_return

    def save_model(self, a_fn, c_fn):
        self.eval_policy.save(a_fn)
        self.eval_critic.save(c_fn)

    def load_model(self, a_fn, c_fn):
        self.eval_policy.load_weights(a_fn)
        self.eval_critic.load_weights(c_fn)
        self.target_critic.load_weights(c_fn)
        print(self.eval_policy.summary())
        print(self.eval_critic.summary())
        print(self.target_critic.summary())

from MAEnv.env_FindGoals.env_FindGoals import EnvFindGoals
from MAEnv.env_Drones.env_Drones import EnvDrones
def run():
    train_steps = 0
    n_epoch = 1000
    n_episodes = 1
    max_episode_len = 100
    env = EnvDrones(map_size=50, drone_num=5, view_range=10, tree_num=30, human_num=20)
    n_actions = 4
    n_agents = 5
    state_shape = 50 * 50 * 3
    obs_shape = (2 * 10 - 1) * (2 * 10 - 1) * 3
    epsilon = 0.8
    agents = COMA(n_actions=n_actions, n_agents=n_agents, state_shape=state_shape, obs_shape=obs_shape)
    for epoch in range(n_epoch):
        episodes = []
        # 收集self.args.n_episodes个episodes
        for episode_idx in range(n_episodes):
            env.reset()
            observations, actions_, rewards, states, actions_onehots, agents_onehots, dones = [], [], [], [], [], [], []
            step = 0
            episode_reward = 0
            last_action = np.zeros((n_agents, n_actions))

            for i in range(max_episode_len):
                #env.render()
                obs = [env.get_drone_obs(env.drone_list[i]).reshape((1, -1))[0] for i in range(len(env.drone_list))]
                state = env.get_joint_obs().reshape(1,-1).repeat(n_agents, axis=0)
                actions, actions_onehot = [], []
                agents_onehot = np.zeros((n_agents,n_agents))
                for agent_id in range(n_agents): #n_agents
                    # 输入当前agent上一个时刻的动作
                    action = agents.choose_action(obs[agent_id], last_action[agent_id],
                                                  agent_id, epsilon=epsilon, evaluate=False)
                    # 生成对应动作的0 1向量
                    action_onehot = np.zeros(n_actions)#n_actions
                    action_onehot[action] = 1

                    agents_onehot[agent_id, agent_id] = 1
                    actions.append(action)
                    actions_onehot.append(action_onehot)
                    last_action[agent_id] = action_onehot

                agents_onehots.append(agents_onehot)

                reward, done = env.drone_step(drone_act_list=actions)

                observations.append(obs)
                states.append(state)
                actions_.append(np.reshape(actions, [n_agents, 1]))#[n_agents, 1]
                actions_onehots.append(actions_onehot)
                rewards.append([reward])

                episode_reward += reward
                step += 1
                if done:
                    dones.append([1])
                    break
                else:
                    dones.append([0])
            # 处理最后一个obs
            observations.append(obs)
            states.append(state)
            observations_next = observations[1:]
            states_next = states[1:]
            observations = observations[:-1]
            states = states[:-1]

            episode = dict(o=observations.copy(),
                           s=states.copy(),
                           u=actions_.copy(),
                           r=rewards.copy(),
                           o_next=observations_next.copy(),
                           s_next=states_next.copy(),
                           u_onehot=actions_onehots.copy(),
                           terminated=dones.copy(),
                           a_onehot = agents_onehots.copy()
                           )
            for key in episode.keys():
                episode[key] = np.array([episode[key]])
            episodes.append(episode)
            print('epoch: {}, episode: {}, episode_reward: {}, episode_steps:{}'.format(epoch, episode_idx, episode_reward, step))
            with agents.summary_writer.as_default():
                tf.summary.scalar('Main/episode_reward', episode_reward, step=epoch)
            agents.summary_writer.flush()
        # episode的每一项都是一个(1, episode_len, n_agents, 具体维度)四维数组，下面要把所有episode的的obs,action等信息拼在一起

        episode_batch = episodes[0]
        episodes.pop(0)
        for episode in episodes:
            for key in episode_batch.keys():
                episode_batch[key] = np.concatenate((episode_batch[key], episode[key]), axis=0)
        #episode_bacth中的每一项(n_episodes, episode_len, n_agents, 具体维度)
        #我们采用最简化设计，n_episodes = 1， 即只采一轮数据就进行训练
        terminated = episode_batch['terminated']
        max_episode_len = terminated.shape[1]
        agents.learn(batch=episode_batch, max_episode_len=max_episode_len, train_step=train_steps, epsilon=epsilon)
        train_steps += 1
    agents.save_model(a_fn="./model/eval_policy.h5", c_fn="./model/eval_critic.h5")

if __name__ == '__main__':
    run()


#!/usr/bin/env python
# -*- coding: utf-8 -*-
# this script uses pretrained model to segment Arabic dialect data.
# it takes the pretrained model trained on joint dialects and the
# training vocab and produces segmented text
#
# Copyright (C) 2017, Qatar Computing Research Institute, HBKU, Qatar
# Las Update: Sun Oct 29 15:34:43 +03 2017
#
# BibTex: @inproceedings{samih2017learning,
#  title={Learning from Relatives: Unified Dialectal Arabic Segmentation},
#  author={Samih, Younes and Eldesouki, Mohamed and Attia, Mohammed and Darwish, Kareem and Abdelali, Ahmed and Mubarak, Hamdy and Kallmeyer, Laura},
#  booktitle={Proceedings of the 21st Conference on Computational Natural Language Learning (CoNLL 2017)},
#  pages={432--441},
#  year={2017}}
#
__author__ = 'Ahmed Abdelali (aabdelali@hbku.edu.qa)'

import tensorflow as tf
import numpy as np
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input
from tensorflow.keras.layers import Dense, Embedding, LSTM, Bidirectional, Dropout
from tensorflow.keras.layers import TimeDistributed
from tensorflow.keras.optimizers import RMSprop
from tensorflow.keras.callbacks import Callback, EarlyStopping, ModelCheckpoint
from tensorflow.keras.models import Sequential
from .ChainCRF import ChainCRF

def build_model(model_path, max_features, word_embedding_dim, maxlen, nb_seg_tags, lstm_dim):
    model = Sequential()
    model.add(Embedding(max_features, word_embedding_dim, input_length=maxlen, name='word_emb', mask_zero=True))
    model.add(Dropout(0.5))
    model.add(Bidirectional(LSTM(lstm_dim,return_sequences=True)))
    model.add(Dropout(0.5))
    model.add(TimeDistributed(Dense(nb_seg_tags)))
    crf = ChainCRF()
    model.add(crf)
    model.compile(loss=crf.sparse_loss,
                   optimizer= RMSprop(0.01),
                   metrics=['sparse_categorical_accuracy'])

    return model