import logging
import pytest
import uuid
from pygame.math import Vector2
import pygame
import random
import numpy as np
import cv2
import argparse
import time

from gobigger.utils import Border
from gobigger.server import Server
from gobigger.render import RealtimeRender, RealtimePartialRender, EnvRender
from gobigger.agents import BotAgent

logging.basicConfig(level=logging.DEBUG)

def play_control_by_keyboard():
    server = Server(dict(
            team_num=1,
            player_num_per_team=1,
            match_time=60*10,
            state_tick_per_second=20, # frame
            action_tick_per_second=5, # frame
        ))
    server.start()
    render = RealtimeRender(server.map_width, server.map_height)
    fps_real = 0
    t1 = time.time()
    clock = pygame.time.Clock()
    fps_set = server.state_tick_per_second
    for i in range(100000):
        clock.tick(1000)
        actions = None
        x, y = None, None
        action_type = -1
        # ================ control by keyboard ===============
        for event in pygame.event.get():
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_UP:
                    x, y = 0, -1
                    action_type = -1
                elif event.key == pygame.K_DOWN:
                    x, y = 0, 1
                    action_type = -1
                elif event.key == pygame.K_LEFT:
                    x, y = -1, 0
                    action_type = -1
                elif event.key == pygame.K_RIGHT:
                    x, y = 1, 0
                    action_type = -1
                elif event.key == pygame.K_q: # Spores
                    x, y = None, None
                    action_type = 0
                elif event.key == pygame.K_w: # Splite
                    x, y = None, None
                    action_type = 1
                elif event.key == pygame.K_e: # Stop moving
                    x, y = None, None
                    action_type = 2
                elif event.key == pygame.K_r: # Cheating, adding weight
                    x, y = None, None
                    action_type = -1
                    server.player_manager.get_players()[0].get_balls()[0].set_size(144)
                actions = {player.name: [x, y, action_type] for player in server.player_manager.get_players()}
        if server.last_time < server.match_time:
            server.step_state_tick(actions=actions)
            if actions is not None and x is not None and y is not None:
                render.fill(server, direction=Vector2(x, y), fps=fps_real, last_time=server.last_time)
            else:
                render.fill(server, direction=None, fps=fps_real, last_time=server.last_time)
            render.show()
            if i % server.state_tick_per_second == 0:
                t2 = time.time()
                fps_real = server.state_tick_per_second/(t2-t1)
                t1 = time.time()
        else:
            logging.debug('Game Over')
            break
        clock.tick(fps_set)
    render.close()

def play_control_by_keyboard_multi_player():
    server = Server(dict(
            team_num=2,
            player_num_per_team=1,
            match_time=60*10,
            state_tick_per_second=20, # frame
            action_tick_per_second=5, # frame
        ))
    server.start()
    render = RealtimeRender(server.map_width, server.map_height)
    player_names = [player.name for player in server.player_manager.get_players()]
    fps_real = 0
    t1 = time.time()
    clock = pygame.time.Clock()
    fps_set = server.state_tick_per_second
    for i in range(100000):
        actions = None
        # ================ control by keyboard ===============
        for event in pygame.event.get():
            if event.type == pygame.KEYDOWN:
                x1, y1, x2, y2 = None, None, None, None
                action_type1 = -1
                action_type2 = -1
                if event.key == pygame.K_UP:
                    x1, y1 = 0, -1
                if event.key == pygame.K_DOWN:
                    x1, y1 = 0, 1
                if event.key == pygame.K_LEFT:
                    x1, y1 = -1, 0
                if event.key == pygame.K_RIGHT:
                    x1, y1 = 1, 0
                if event.key == pygame.K_LEFTBRACKET: # Spores
                    action_type1 = 0
                if event.key == pygame.K_RIGHTBRACKET: # Splite
                    action_type1 = 1
                if event.key == pygame.K_BACKSLASH: # Stop moving
                    action_type1 = 2
                if event.key == pygame.K_w:
                    x2, y2 = 0, -1
                if event.key == pygame.K_s:
                    x2, y2 = 0, 1
                if event.key == pygame.K_a:
                    x2, y2 = -1, 0
                if event.key == pygame.K_d:
                    x2, y2 = 1, 0
                if event.key == pygame.K_1: # Spores
                    action_type2 = 0
                if event.key == pygame.K_2: # Splite
                    action_type2 = 1
                if event.key == pygame.K_3: # Stop moving
                    action_type2 = 2
                actions = {
                    player_names[0]: [x1, y1, action_type1],
                    player_names[1]: [x2, y2, action_type2],
                }
        if server.last_time < server.match_time:
            server.step_state_tick(actions=actions)
            if actions is not None and x1 is not None and y1 is not None:
                render.fill(server, direction=Vector2(x1, y1), fps=fps_real, last_time=server.last_time)
            else:
                render.fill(server, direction=None, fps=fps_real, last_time=server.last_time)
            render.show()
            if i % server.state_tick_per_second == 0:
                t2 = time.time()
                fps_real = server.state_tick_per_second/(t2-t1)
                t1 = time.time()
        else:
            logging.debug('Game Over')
            break
        clock.tick(fps_set)
    render.close()

def play_control_by_keyboard_partial():
    '''
    Keyboard control + partly visible
    '''
    server = Server(dict(
            team_num=1,
            player_num_per_team=1,
            match_time=60*10,
            state_tick_per_second=20, # frame
            action_tick_per_second=5, # frame
        ))
    server.start()
    render = RealtimePartialRender(server.map_width, server.map_height)
    fps_real = 0
    t1 = time.time()
    clock = pygame.time.Clock()
    fps_set = server.state_tick_per_second
    for i in range(100000):
        actions = None
        x, y = None, None
        action_type = -1
        # ================ control by keyboard ===============
        for event in pygame.event.get():
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_UP:
                    x, y = 0, -1
                    action_type = -1
                elif event.key == pygame.K_DOWN:
                    x, y = 0, 1
                    action_type = -1
                elif event.key == pygame.K_LEFT:
                    x, y = -1, 0
                    action_type = -1
                elif event.key == pygame.K_RIGHT:
                    x, y = 1, 0
                    action_type = -1
                elif event.key == pygame.K_q: # Spore
                    x, y = None, None
                    action_type = 0
                elif event.key == pygame.K_w: # Splite
                    x, y = None, None
                    action_type = 1
                elif event.key == pygame.K_e: # stop moving
                    x, y = None, None
                    action_type = 2
                elif event.key == pygame.K_r: # cheating
                    x, y = None, None
                    action_type = -1
                    server.player_manager.get_players()[0].get_balls()[0].set_size(144)
                actions = {player.name: [x, y, action_type] for player in server.player_manager.get_players()}

        if server.last_time < server.match_time:
            server.step_state_tick(actions=actions)
            if actions is not None and x is not None and y is not None:
                render.fill(server, direction=Vector2(x, y), fps=fps_real, last_time=server.last_time)
            else:
                render.fill(server, direction=None, fps=fps_real, last_time=server.last_time)
            render.show()
            if i % server.state_tick_per_second == 0:
                t2 = time.time()
                fps_real = server.state_tick_per_second/(t2-t1)
                t1 = time.time()
        else:
            logging.debug('Game Over')
            break
        clock.tick(fps_set)
    render.close()

def play_control_by_keyboard_vs_bot(team_num=2):
    server = Server(dict(
            team_num=team_num,
            player_num_per_team=1,
            match_time=60*10,
            save_video=True,
            state_tick_per_second=20, # frame
            action_tick_per_second=5, # frame
        ))
    server.start()
    render = RealtimeRender(server.map_width, server.map_height)
    server.set_render(render)
    human_team_name = '0'
    human_team_player_name = []
    bot_agents = []
    for player in server.player_manager.get_players():
        if player.team_name != human_team_name:
            bot_agents.append(BotAgent(player.name))
        else:
            human_team_player_name.append(player.name)
    fps_real = 0
    t1 = time.time()
    clock = pygame.time.Clock()
    fps_set = server.state_tick_per_second
    for i in range(100000):
        obs = server.obs()
        actions_bot = {bot_agent.name: bot_agent.step(obs[1][bot_agent.name]) for bot_agent in bot_agents}
        actions = {player_name: [None, None, -1] for player_name in human_team_player_name}
        x, y = None, None
        action_type = -1
        # ================ control by keyboard ===============
        for event in pygame.event.get():
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_UP:
                    x, y = 0, -1
                    action_type = -1
                elif event.key == pygame.K_DOWN:
                    x, y = 0, 1
                    action_type = -1
                elif event.key == pygame.K_LEFT:
                    x, y = -1, 0
                    action_type = -1
                elif event.key == pygame.K_RIGHT:
                    x, y = 1, 0
                    action_type = -1
                elif event.key == pygame.K_q: # Spores
                    x, y = None, None
                    action_type = 0
                elif event.key == pygame.K_w: # Splite
                    x, y = None, None
                    action_type = 1
                elif event.key == pygame.K_e: # Stop moving
                    x, y = None, None
                    action_type = 2
                elif event.key == pygame.K_r: # Cheating, adding weight
                    x, y = None, None
                    action_type = -1
                    server.player_manager.get_players()[0].get_balls()[0].set_size(144)
                actions = {player_name: [x, y, action_type] for player_name in human_team_player_name}
        if server.last_time < server.match_time:
            actions.update(actions_bot)
            server.step_state_tick(actions=actions)
            if actions is not None and x is not None and y is not None:
                render.fill(server, direction=Vector2(x, y), fps=fps_real, last_time=server.last_time)
            else:
                render.fill(server, direction=None, fps=fps_real, last_time=server.last_time)
            render.show()
            if i % server.state_tick_per_second == 0:
                t2 = time.time()
                fps_real = server.state_tick_per_second/(t2-t1)
                t1 = time.time()
        else:
            logging.debug('Game Over')
            break
        clock.tick(fps_set)
    render.close()

def play_by_bot(team_num, player_num, agent_class):

    class TeamAgent:
        def __init__(self, team_name, player_names):
            self.agents = {}
            self.team_name = team_name
            self.player_names = player_names
            for player_name in self.player_names:
                self.agents[player_name] = BotAgent(name=player_name)

        def get_actions(self, obs):
            global_state, player_states = obs
            actions = {}
            for player_name, agent in self.agents.items():
                action = agent.step(player_states[player_name])
                actions[player_name] = action
            return actions

    server = Server(dict(
        team_num=team_num,
        player_num_per_team=player_num,  # 每个队伍的玩家数量
        match_time=60 * 10,  # 每场比赛的持续时间
        save_video=False,
        state_tick_per_second=20, # frame
        action_tick_per_second=5, # frame
    ))
    render = RealtimeRender(server.map_width, server.map_height)
    server.set_render(render)
    server.start()
    agents = []
    team_player_names = server.get_team_names()
    team_names = list(team_player_names.keys())
    for index in range(server.team_num):
        try:
            if agent_class == '':
                p = TeamAgent
            else:
                from importlib import import_module
                class_path, class_name = agent_class.split(':')
                p = import_module(class_path)
                p = getattr(p, class_name)
            agents.append(p(team_name=team_names[index],
                                         player_names=team_player_names[team_names[index]]))
        except Exception as e:
            print('load bot class failed.')
            exit()

    for i in range(60 * 10 * server.action_tick_per_second + 1):
        obs = server.obs()
        render.show()
        event = pygame.event.poll()
        render.fill(server, direction=None, fps=60, last_time=server.last_time)
        global_state, player_states = obs
        actions = {}
        for agent in agents:
            agent_obs = [global_state, {
                player_name: player_states[player_name] for player_name in agent.player_names
            }]
            actions.update(agent.get_actions(agent_obs))
        finish_flag = server.step(actions=actions)
        if finish_flag:
            logging.debug('Game Over!')
            break
    server.close()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--player-num', type=int, choices=[1,2,3], default=1)
    parser.add_argument('--vision-type', type=str, choices=['full', 'partial'], default='full')
    parser.add_argument('--vs-bot', action='store_true', help='vs bot')
    parser.add_argument('--bot-only', action='store_true', help='bot only')
    parser.add_argument('--team-num', type=int, default=2)
    parser.add_argument('--agent-class', type=str, default='', help='agent class, like \''
                                                                    'submit.bot_submission:BotSubmission\'')
    args = parser.parse_args()

    if args.bot_only:
        play_by_bot(team_num=args.team_num, player_num=args.player_num, agent_class=args.agent_class)
    elif args.vs_bot:
        play_control_by_keyboard_vs_bot(team_num=args.team_num)
    elif args.player_num == 1 and args.vision_type == 'full':
        play_control_by_keyboard()
    elif args.player_num == 1 and args.vision_type == 'partial':
        play_control_by_keyboard_partial()
    elif args.player_num == 2 and args.vision_type == 'full':
        play_control_by_keyboard_multi_player()
    else:
        logging.error('Not supoort when player_num = {} and vision_type = {}'.format(args.player_num, args.vision_type))


import argparse
import os
import numpy as np
import math
import itertools
import time
import datetime
import sys

from torchvision.utils import save_image

from torch.utils.data import DataLoader
from torchvision import datasets

from models import *
from datasets import *

import torch.nn as nn
import torch.nn.functional as F
import torch

from scipy import misc

parser = argparse.ArgumentParser()
parser.add_argument('--epoch', type=int, default=0, help='epoch to start training from')
parser.add_argument('--n_epochs', type=int, default=200, help='number of epochs of training')
parser.add_argument('--dataset_name', type=str, default="facades", help='name of the dataset')
parser.add_argument('--batch_size', type=int, default=1, help='size of the batches')
parser.add_argument('--lr', type=float, default=0.0002, help='adam: learning rate')
parser.add_argument('--b1', type=float, default=0.5, help='adam: decay of first order momentum of gradient')
parser.add_argument('--b2', type=float, default=0.999, help='adam: decay of first order momentum of gradient')
parser.add_argument('--decay_epoch', type=int, default=100, help='epoch from which to start lr decay')
parser.add_argument('--n_cpu', type=int, default=8, help='number of cpu threads to use during batch generation')
parser.add_argument('--img_height', type=int, default=256, help='size of image height')
parser.add_argument('--img_width', type=int, default=256, help='size of image width')
parser.add_argument('--channels', type=int, default=3, help='number of image channels')
parser.add_argument('--sample_interval', type=int, default=500, help='interval between sampling of images from generators')
parser.add_argument('--checkpoint_interval', type=int, default=-1, help='interval between model checkpoints')
opt = parser.parse_args()
print(opt)

os.makedirs('images/%s' % opt.dataset_name, exist_ok=True)
os.makedirs('saved_models/%s' % opt.dataset_name, exist_ok=True)

cuda = True if torch.cuda.is_available() else False
device = torch.device("cuda" if cuda else "cpu")

# Loss functions
criterion_GAN = torch.nn.MSELoss()
#criterion_pixelwise = torch.nn.L1Loss()
criterion_pixelwise = torch.nn.MSELoss()

# Loss weight of L1 pixel-wise loss between translated image and real image
lambda_pixel = 100

# Initialize generator and discriminator
generator = GeneratorUNet(in_channels=2, out_channels=2).to(device)
#discriminator = Discriminator().to(device)

if cuda:
    criterion_GAN.cuda()
    criterion_pixelwise.cuda()

# Load pretrained models
generator.load_state_dict(torch.load('saved_models/facades/generator_end.pth'))
#discriminator.load_state_dict(torch.load('saved_models/%s/discriminator_%d.pth' % (opt.dataset_name, opt.epoch)))

# Configure dataloaders
dataloader = DataLoader(Radars(),batch_size=opt.batch_size, shuffle=True, num_workers=1)

# ----------
#  Training
# ----------

prev_time = time.time()

num = 0
for epoch in range(opt.epoch, opt.n_epochs):
    for i, batch in enumerate(dataloader):
        A, B = batch
        A = A.to(device)
        B = B.to(device)

        # Model inputs

        # ------------------
        #  Train Generators
        # ------------------

        # GAN loss
        fake_B = generator(A)
        m = fake_B.to(torch.device('cpu')).detach().numpy()
        n = A.to(torch.device('cpu')).detach().numpy()
        l = B.to(torch.device('cpu')).detach().numpy()
        res = np.concatenate((m[0],l[0],n[0]))
        print(res.shape)
        print(np.mean(res))
        misc.imsave('images/music/%d_%d.png' % (epoch,i),np.concatenate((m[0],l[0],n[0])))
        loss_pixel = criterion_pixelwise(fake_B, B)
        num += 1
        if num >100:
            break

        # --------------
        #  Log Progress
        # --------------

        # Determine approximate time left
        batches_done = epoch * len(dataloader) + i
        batches_left = opt.n_epochs * len(dataloader) - batches_done
        time_left = datetime.timedelta(seconds=batches_left * (time.time() - prev_time))
        prev_time = time.time()

        # Print log
        sys.stdout.write("\r[Epoch %d/%d] [Batch %d/%d] [pixel: %f ] ETA: %s" %
                                                        (epoch, opt.n_epochs,
                                                        i, len(dataloader),
                                                        loss_pixel.item(),
                                                        time_left))

#        if batches_done % opt.sample_interval == 0:
#            sample_images(batches_done)


import os
import numpy as np
import json
import logging
import random
import urllib.request
from urllib.parse import unquote

from pipeline_monitor import prometheus_monitor as monitor
from pipeline_logger import log

import tensorflow as tf
from tensorflow.contrib import predictor

from tensorflow.keras.applications.mobilenet import preprocess_input, decode_predictions
from tensorflow.keras.preprocessing import image
import glob

_logger = logging.getLogger('pipeline-logger')
_logger.setLevel(logging.INFO)
_logger_stream_handler = logging.StreamHandler()
_logger_stream_handler.setLevel(logging.INFO)
_logger.addHandler(_logger_stream_handler)

__all__ = ['invoke']

_labels = {
           'name': 'transfer',
           'tag': 'v1',
           'runtime': 'python',
           'chip': 'cpu',
           'resource_type': 'model',
           'resource_subtype': 'keras',
          }

def _initialize_upon_import():
    """Initialize / Restore Model Object."""

    saved_model_list = sorted(glob.glob("./pipeline_tfserving/*"))

    #saved_model_path = './pipeline_tfserving/0'
    saved_model_path = saved_model_list[-1]
    print('Loading model found at %s'  % saved_model_path)
    loaded_model = tf.contrib.saved_model.load_keras_model(saved_model_path)
    return loaded_model

_classes_list = sorted(glob.glob("./images/predict/*"))

# This is called unconditionally at *module import time*...
_model = _initialize_upon_import()

@log(labels=_labels, logger=_logger)
def invoke(request):
    """Where the magic happens..."""

    with monitor(labels=_labels, name="transform_request"):
        transformed_request = _transform_request(request)

    with monitor(labels=_labels, name="invoke"):
        response = _model.predict(transformed_request)

    with monitor(labels=_labels, name="transform_response"):
        transformed_response = _transform_response(response)

    return transformed_response

def download_image(url):
    name = random.randrange(1,100000)
    fullname = '/tmp/' + str(name) + ".img"
    urllib.request.urlretrieve(url, fullname)
    return fullname

def _transform_request(request):
    request = request.decode('utf-8')
    request = unquote(request)

    # Direct http example
    if request.startswith('http'):
        request = download_image(request)
    else:
        # Slack Label Example
        request_array = request.split('&')
        print(request_array)

        result = [value for value in request_array if value.startswith('text=')]
        if len(result) > 0:
            request = download_image(result[0][5:])
            print(request)

    predict_img = image.load_img(request, target_size=(224, 224))
    predict_img_array = image.img_to_array(predict_img)
    predict_img_array = np.expand_dims(predict_img_array, axis=0)
    predict_preprocess_img = preprocess_input(predict_img_array)

    return predict_preprocess_img

def _transform_response(response):

    # response_formatted = '\\n '.join("%s%%\t%s" % (str((100 * item['score']))[0:4], item['name']) for item in classification_response_json)

    return json.dumps({"classes": [item.split('/')[-1] for item in _classes_list],
                       "probabilities": ['%.2f%%' % (item * 100) for item in response.tolist()[0]],
                      })

if __name__ == '__main__':
#    request = './images/predict/cat.jpg'
#    request = b'http%3A%2F%2Fsite.meishij.net%2Fr%2F58%2F25%2F3568808%2Fa3568808_142682562777944.jpg'

    request = b'token=asdasd&team_id=T6QHWMRD4&team_domain=pipelineai&channel_id=adfasdf&channel_name=privategroup&user_id=asdfasdf&user_name=cfregly&command=/predict&text=https://images.ctfassets.net/kvimhx6nhg7h/5WclEHFxUksuS2IwsUECE6/a29fa96920666f9d4eb7c456403e4f9d/Tan-cat-in-a-cone.png&response_url=https://hooks.slack.com/commands/T/5/r&trigger_id=3d3'

    response = invoke(request)
    print(response)
#    with open('./pipeline_test_request.json', 'rb') as fb:
#        request_bytes = fb.read()
#        response_bytes = invoke(request_bytes)
#        print(response_bytes)


# ------------------------------------------------------------------------------
# Program:     The LDAR Simulator (LDAR-Sim)
# File:        methods.company
# Purpose:     Initialize companies, generate crews, deploy crews and reporting
#
# Copyright (C) 2018-2021  Intelligent Methane Monitoring and Management System (IM3S) Group
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the MIT License as published
# by the Free Software Foundation, version 3.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# MIT License for more details.

# You should have received a copy of the MIT License
# along with this program.  If not, see <https://opensource.org/licenses/MIT>.
#
# ------------------------------------------------------------------------------

import math
from datetime import timedelta
from importlib import import_module

import numpy as np
from methods.deployment.generic_funcs import get_deployment_dates
from utils.attribution import update_flag
from utils.generic_functions import get_prop_rate

class BaseCompany:
    """ Base companies are used by methods to generate crews, company-level scheduling,
        deploy crews, flag sites and carry out reporting.
    """

    def __init__(self, state, parameters, config, timeseries, module_name):
        """
        Initialize a company to manage the crews (e.g. a contracting company).

        """
        self.name, self.state, = config['label'], state
        self.parameters = parameters
        self.config = config
        self.timeseries = timeseries
        self.site_watchlist = {}
        self.crews = []
        deploy_mod = import_module('methods.deployment.{}_company'.format(
            self.config['deployment_type'].lower()))
        Schedule = getattr(deploy_mod, 'Schedule')
        self.schedule = Schedule(config, parameters, state)
        self.deployment_years, self.deployment_months = get_deployment_dates(config, state)
        self.deployment_days = self.state['weather'].deployment_days(
            method_name=self.name,
            config=config,
            start_date=self.state['t'].start_date,
            start_work_hour=8,  # Start hour in day
            consider_weather=parameters['consider_weather'])

        # --- init time series ---
        n_ts = parameters['timesteps']
        self.timeseries['{}_prop_sites_avail'.format(self.name)] = np.zeros(n_ts)
        self.timeseries['{}_cost'.format(self.name)] = np.zeros(n_ts)
        self.timeseries['{}_sites_visited'.format(self.name)] = np.zeros(n_ts)
        self.timeseries['{}_travel_time'.format(self.name)] = np.zeros(n_ts)
        self.timeseries['{}_survey_time'.format(self.name)] = np.zeros(n_ts)
        self.timeseries['{}_redund_tags'.format(self.name)] = np.zeros(n_ts)
        self.timeseries['{}_missed_leaks'.format(self.name)] = np.zeros(n_ts)
        self.timeseries['{}_sites_vis_w_leaks'.format(self.name)] = np.zeros(n_ts)
        self.timeseries['{}_eff_flags'.format(self.name)] = np.zeros(n_ts)
        self.timeseries['{}_n_tags'.format(self.name)] = np.zeros(n_ts)
        self.timeseries['{}_flags_redund1'.format(self.name)] = np.zeros(n_ts)
        self.timeseries['{}_flags_redund2'.format(self.name)] = np.zeros(n_ts)
        self.timeseries['{}_flag_wo_vent'.format(self.name)] = np.zeros(n_ts)

        if self.config["measurement_scale"] != 'component':
            # Assign the correct follow-up threshold
            if self.config['follow_up']['threshold_type'] == "absolute":
                self.config['follow_up']['thresh'] = self.config['follow_up']['threshold']
            elif self.config['follow_up']['threshold_type'] == "proportion":
                self.config['follow_up']['thresh'] = get_prop_rate(
                    self.config['follow_up']['proportion'],
                    self.state['empirical_leaks'])
            else:
                print('Follow-up thresh type not recognized. Must be "absolute" or "proportion".')

        # ---Init 2D matrices to store deployment day (DD) counts and MCBs ---
        self.DD_map = np.zeros(
            (len(self.state['weather'].longitude),
             len(self.state['weather'].latitude)))
        self.MCB_map = np.zeros(
            (len(self.state['weather'].longitude),
             len(self.state['weather'].latitude)))

        # --- init site specific variables ---
        for site in self.state['sites']:
            survey_min_interval = '{}_min_int'.format(self.name)
            if survey_min_interval in site:
                t_since_last_LDAR = site[survey_min_interval]
            else:
                t_since_last_LDAR = 0
            site.update({'{}_t_since_last_LDAR'.format(self.name): t_since_last_LDAR})
            site.update({'{}_surveys_conducted'.format(self.name): 0})
            site.update({'{}_attempted_today?'.format(self.name): False})
            site.update({'{}_surveys_done_this_year'.format(self.name): 0})
            site.update({'{}_missed_leaks'.format(self.name): 0})

        make_crew_loc = import_module('methods.deployment.{}_company'.format(
            self.config['deployment_type'].lower()))
        make_crews = getattr(make_crew_loc, 'make_crews')
        make_crews(self.crews, config, state, parameters, timeseries, self.deployment_days)
        for idx, cnt in enumerate(self.crews):
            self.timeseries['{}_cost'.format(self.name)][self.state['t'].current_timestep] += \
                self.config['cost']['upfront']

        self.schedule.assign_agents()

    def deploy_crews(self):
        """
        The company tells all the crews to get to work.
        """
        # NOTE: crew checks weather conditions at start of the day
        c_date = self.state['t'].current_date
        if c_date.month in self.deployment_months and c_date.year in self.deployment_years:
            self.candidate_flags = []
            if self.config['is_follow_up']:
                site_pool = self.state['flags']
            else:
                site_pool = self.state['sites']
            # Get sites that are ready for survey
            site_pool = self.schedule.get_due_sites(site_pool)
            # Get number of crews working that day based on number of sites ready for visit
            if self.config['deployment_type'] == 'stationary':
                # assume all crews/sensors are working
                n_working_crews = len(self.crews)
            else:
                n_working_crews = self.schedule.get_working_crews(site_pool, self.config['n_crews'])
            for idx in range(n_working_crews):
                # Triage sites to crew
                crew_site_list = self.schedule.get_crew_site_list(site_pool, idx,
                                                                  n_working_crews, self.crews)
                # Send crew to site
                if len(crew_site_list) > 0:
                    self.crews[idx].work_a_day(crew_site_list, self.candidate_flags)
            if len(self.candidate_flags) > 0:
                self.candidates_to_watchlist()

            # Calculate proportion sites available based on weather of site
            available_sites = sum([1 for site in self.state['sites']
                                   if self.deployment_days[site['lon_index'], site['lat_index'],
                                                           self.state['t'].current_timestep]])
            prop_avail = available_sites / len(self.state['sites'])
            self.timeseries['{}_prop_sites_avail'.format(
                self.name)][self.state['t'].current_timestep] = prop_avail
        else:
            self.timeseries['{}_prop_sites_avail'.format(
                self.name)][self.state['t'].current_timestep] = 0

        # Update time series stats

        days_since_LDAR = '{}_t_since_last_LDAR'.format(self.name)
        attempted_today = '{}_attempted_today?'.format(self.name)
        sites_this_year = '{}_surveys_done_this_year'.format(self.name)
        for site in self.state['sites']:
            site[days_since_LDAR] += 1
            site[attempted_today] = False
            if self.state['t'].current_date.day == 1 \
                    and self.state['t'].current_date.month == 1:
                site[sites_this_year] = 0

        # Daily check of watchlist
        if len(self.site_watchlist) > 0:
            self.flag_sites()

        # Update followup specific parameters
        if self.config['is_follow_up']:
            self.state['flags'] = [site['facility_ID'] for site in self.state['sites']
                                   if site['currently_flagged']]
        return

    def candidates_to_watchlist(self):
        """
        Add candidate sites to watchlist, and calculate a site emission rate if there are
        multiple emissions measurements.The watchlist allows companies to keep track of
        canditate site surveys. This is utilized by a company when a follow_up_delay
        is greater than zero, as multiple surveys can be done for each site before a site
        is flagged.

        If an instant repair threshold is utilized and a sites  emissions rate is greater than
        the instant followup threshold, the site will be  immediately flagged.

        """
        site_wl = self.site_watchlist

        # Go through Each Candidate site, and add it to the watchlist
        # A watchlist keeps track of sites with leaks that have not yet been flagged
        for site in self.candidate_flags:
            facility_ID = site['site']['facility_ID']
            # Get instant follow-up sites and init follow-up sites based on follow_up delay
            if self.config['follow_up']['instant_threshold'] \
                    and site['site_measured_rate'] > self.config['follow_up']['instant_threshold']:
                self.flag_site(site)
                # Remove the site from the watchlist if it needs an instant followup
                if facility_ID in site_wl:
                    site_wl.pop(facility_ID)
            elif facility_ID not in site_wl:
                site_wl.update(
                    {facility_ID: {
                        'site': site['site'],
                        'dates': [self.state['t'].current_date],
                        'measured_emis_rates': [site['site_measured_rate']],
                        'measured_vent_rates': [site['vent_rate']],
                        'site_measured_rate': site['site_measured_rate'],
                        'vent_rate': site['vent_rate']
                    }})
            else:
                site_wl[facility_ID]['dates'].append(self.state['t'].current_date)
                site_wl[facility_ID]['measured_emis_rates'].append(site['site_measured_rate'])
                site_wl[facility_ID]['measured_vent_rates'].append(site['vent_rate'])
                # Calculate effective emissions and vent rates
                site_wl[facility_ID]['site_measured_rate'] = self.aggregate_by(
                    site_wl[facility_ID]['measured_emis_rates'],
                    self.config['follow_up']['redundancy_filter']
                )
                site_wl[facility_ID]['vent_rate'] = self.aggregate_by(
                    site_wl[facility_ID]['measured_vent_rates'],
                    self.config['follow_up']['redundancy_filter']
                )

        # Sort Watchlist by emissions rate
        # This ensures that the leakest sites are at the top of the dictionary
        # so proportional site followup can be performed without resorting everyday
        # within the flag_sites function.
        site_wl = {k: v for k, v in sorted(
            site_wl.items(),
            key=lambda x: x[1]['site_measured_rate'],
            reverse=True)
        }

    def flag_site(self, site):
        """ Flag a single site and check for leak redundancy

        Args:
            site (dict): a dictionary which must have a site style dict,
                         (state['sites']) a site_measured_rate, and a site
                         vent rate
        """
        self.site_watchlist.pop(site['site']['facility_ID'], None)
        update_flag(self.config, site, self.timeseries, self.state['t'],
                    self.name, self.parameters['emissions']['consider_venting'])

    def flag_sites(self):
        """
        Performs followup proportion and threshold checks on all sites for watchlist sites
        past follow-up delay / grace period.

        """
        site_wl = self.site_watchlist
        first_survey_time = min(site['dates'][0] for sidx, site in site_wl.items())
        can_flag_date = first_survey_time + timedelta(days=self.config['follow_up']['delay'])
        # If the delay / grace period has passed since the first leak reported in the watchlist
        # start surveys. Date is used to ignore the hour Leak was detected and focus on the day
        # instead.
        if self.state['t'].current_date.date() >= can_flag_date.date():
            # Get portion of sites and threshold filtered sited
            if self.config['follow_up']['interaction_priority'] == 'proportion':
                prop_to_flag = int(math.ceil(
                    len(site_wl) * self.config['follow_up']['proportion']))
                # site_wl is an ordered by leak size.
                site_wl_subset = site_wl.items()[:prop_to_flag]
                target_sites = [site for sdx, site in site_wl_subset
                                if site['site_measured_rate']
                                > self.config['follow_up']['threshold']]
            if self.config['follow_up']['interaction_priority'] == 'threshold':
                sites_above_thresh = [site for sdx, site in site_wl.items()
                                      if site['site_measured_rate']
                                      > self.config['follow_up']['threshold']]
                prop_to_flag = int(math.ceil(
                    len(sites_above_thresh) * self.config['follow_up']['proportion']))
                target_sites = sites_above_thresh[:prop_to_flag]
            # Go through all targeted sites and flag them or add them to redundant lists
            for targ_site in target_sites:
                self.flag_site(targ_site)
            # Reset watchlist
            self.site_watchlist = {}

    def aggregate_by(self, list, filter):
        '''
        Gets the affective aggregated value from a list based on an input filter
        Values of the filter can be 'recent', 'max', or 'mean'

        '''
        if filter == 'recent':
            return list[-1]
        elif filter == 'max':
            return max(list)
        elif filter == 'mean':
            return np.average(list)

    def site_reports(self):
        """
        Writes site-level deployment days (DDs) and maximum condition blackouts (MCBs)
        for each site.
        """

        for site in self.state['sites']:
            site['{}_prop_DDs'.format(
                self.name)] = self.DD_map[site['lon_index'], site['lat_index']]
            site['{}_MCB'.format(
                self.name)] = self.MCB_map[site['lon_index'], site['lat_index']]

        return


import tensorflow as tf
import numpy as np
import math
import sys
import glob
import os
import time
np.set_printoptions(threshold = 1e6)
from SpeechSynthesizer import speech_synthesizer
from PhoneClassifer import phone_classifer
from utils.load_data import get_data_ss,default_config,phns
from LOAD_SAVE_PARAMS.LOAD_SAVE_PARAMS import save_weight,load_weight

from tensorflow.compat.v1 import ConfigProto
from tensorflow.compat.v1 import InteractiveSession

config = ConfigProto()

config.gpu_options.allow_growth = True
session = InteractiveSession(config=config)
tf.compat.v1.keras.backend.set_session(session)
tf.compat.v1.keras.backend.clear_session()

os.environ['CUDA_VISIBLE_DEVICES'] = '/gpu:0'

def train(model,params,batch_size):
    wav_files_list = glob.glob(default_config["ss_pre_train"])
    print(len(wav_files_list))
    iter_data=get_data_ss(wav_files_list,batch_size,random_crop=True,bReSample=True)
    los=[]
    count=1
    for x_mfccs,y_spec,y_mel in iter_data:
        e_start=time.time()

        X_mfccs=tf.concat(x_mfccs,0)

        Y_spec=tf.concat(y_spec,0)

        Y_mel=tf.concat(y_mel,0)

        X_mfccs=tf.cast(X_mfccs,dtype=tf.float32)

        Y_spec=tf.cast(Y_spec,dtype=tf.float32)

        Y_mel=tf.cast(Y_mel,dtype=tf.float32)

        with tf.GradientTape() as tape:
            tape.watch(params)

            pred_mel,pred_spec=model(X_mfccs,dropout_rate=0.0)

            loss=model.loss_ss(pred_mel,Y_mel,pred_spec,Y_spec)

        print("loss:%f"%loss)
        los.append(loss)
        grads=tape.gradient(loss,params)
        #grads,_tf.clip_by_global_norm(grads,clip_norm)
        model.update_params(grads,params)
        e_end=time.time()
        print("一次耗时%f秒"%(e_end-e_start))
        if count%1000==0:
            save_weight("ckp","params_ss_pre_train",params)

            filepath="loss_ss_pre_train.txt"
            flie=open(filepath,"a+")
            flie.write(str(tf.math.reduce_mean(los).numpy())+"\n")
            flie.close()

            count=count+1
        else:
            count=count+1

    filepath="loss_ss_pre_train.txt"
    flie=open(filepath,"a+")

    flie.write(str(tf.math.reduce_mean(los).numpy())+"\n")
    flie.close()

def return_accuracy(Y,Y_pre):
    #print(Y.shape)
    #print(Y_pre.shape)
    #print(Y)
    #print(Y_pre)
    num=Y.shape[0]*Y.shape[1]

    rowMaxSoft=np.argmax(Y_pre, axis=-1)+1
    rowMax=np.argmax(Y, axis=-1)+1

    rowMaxSoft=rowMaxSoft.reshape([1,-1])
    rowMax=rowMax.reshape([1,-1])

    nonO=rowMaxSoft-rowMax

    exist = (nonO != 0) * 1.0
    factor = np.ones([nonO.shape[1],1])
    res = np.dot(exist, factor)

    accuracy=(float(num)-res[0][0])/float(num)

    return accuracy

if __name__ == "__main__":

    batch_size=1

    n_mel=default_config["n_mels"]

    n_spec=(default_config["n_fft"]//2)+1

    input_nums=default_config["n_mfcc"]

    num_hiddens=512#default_config["hidden_units"]#768#768电脑带不动

    num_outputs=default_config["n_mfcc"]

    layer_nums=12

    multi_head=12

    mel_layers=6

    spec_layers=6

    max_position_dim=1024

    clip_norm=1.0

    epochs=3000

    isContinue=True

    #lr=1e-5#5e-6
    model_pc=phone_classifer(lr=1e-5,input_nums=input_nums,hidden_nums=num_hiddens,output_nums=num_outputs,max_position_dim=max_position_dim,layers_encoder=layer_nums,labels_num=len(phns),multi_head=multi_head)

    params_pc=model_pc.get_params()+model_pc.get_params_vc()

    model_ss=speech_synthesizer(lr=1e-5,mel_in_nums=len(phns),mel_hi_nums=num_hiddens,mel_out_nums=len(phns),mel_layers=mel_layers,n_mel=n_mel,spec_in_nums=n_mel,spec_hi_nums=num_hiddens,spec_out_nums=n_mel,spec_layers=spec_layers,n_spec=n_spec,max_position_dim=max_position_dim,multi_head=multi_head)

    params_mel=model_ss.get_params_mel()

    params_spec=model_ss.get_params_spec()

    params_ss=params_mel+params_spec

    try:
        load_weight("ckp","params_pc",params_pc)
    except:
        raise("未发现已经训练好的音素分类模型数据")

    model_ss.init_pc(model_pc)#加载pc模型

    if isContinue==True:
        try:
            load_weight("ckp","params_ss",params_ss)

        except:
            raise("未发现已经训练好的语音合成模型数据")

    for i in range(epochs):
        with tf.device('/gpu:0'):
            train(model_ss,params_ss,batch_size)
            save_weight("ckp","params_ss_pre_train",params_ss)