# -*- coding: utf-8 -*-
"""
Created on Tue Dec 10 11:00:31 2019

@author: cimlab
"""

import os
import time
import simpy
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors

np.random.seed(999)

#entity
class Job:
    def __init__(self, ID, jtype, AT, DD, routing, PT):
        self.ID    = ID
        self.jtype = jtype
        self.AT    = AT    #AT: arrival time
        self.DD    = DD    #DD: due date
        self.CT    = None    #CT: complete time

        self.PT       = PT
        self.routing  = routing
        self.progress = 0

#resource in factory
class Source:
    def __init__(self, fac, AV1, AV2, AV3, DD_factor):
        self.fac = fac
        #attribute
        self.AV_interval = {"a": AV1, "b": AV2, "c": AV3}
        self.DD_factor = DD_factor
        self.df = pd.DataFrame(columns = ["ID", "jtype", "arrival_interval", "due_date", "routing", 'PT'])

    def initialize(self):#, OP):
        #reference
        self.env = self.fac.env
        self.queues = self.fac.queues
        # self.OP = OP    #OP: output port, connect to a queue
        #statistics
        self.output = 0
        #initial process
        self.process_A = self.env.process(self.arrival("a"))
        self.process_B = self.env.process(self.arrival("b"))
        self.process_C = self.env.process(self.arrival("c"))

    def arrival(self, jtype):
        while True:
            #wait until job arrive
            arrival_interval = self.AV_interval[jtype]
            arrival_interval = np.random.exponential(arrival_interval)
            yield self.env.timeout(arrival_interval)
            #generate job
            ID      = self.output
            routing = self.fac.routing_dict[jtype].split(',')
            PT      = [int(i) for i in self.fac.PT_dict[jtype].split(',')]
            AT      = self.env.now
            DD      = AT + self.DD_factor * self.fac.avg_pt[jtype]
            job     = Job(ID, jtype, AT, DD, routing, PT)
            # job.routing = routing#.split(',')
            # job.PT      = PT
            self.df.loc[ID] = [ID, jtype, arrival_interval, DD, routing, PT]

            if self.fac.log:
                print("{} job {} release.".format(self.env.now, job.ID))

            self.output += 1
            #send job to next station
            target = job.routing[job.progress]
            self.queues[target].job_arrive(job)
            # self.OP.job_arrive(job)
            self.fac.update_WIP(1)
            '''
            event_trace
            '''
            self.fac.print_state("arrival")

class Queue:
    def __init__(self, fac, ID, DP_rule):
        self.fac = fac
        self.stage = ID
        self.DP_rule = DP_rule    #DP_rule: dispatching rule

    def initialize(self):
        #reference
        self.env = self.fac.env
        self.OP = [self.fac.machines[self.stage]]
        #attribute
        self.space = []
        #statistics
        #initial process

    def MST(self, MC):
        if MC.mode == None:
            job = np.random.choice(self.space)
        else:
            ST_list = []
            for j in self.space:
                _from = MC.mode.jtype
                to = j.jtype
                if _from == to:
                    ST_list.append(0)
                else:
                    ST = self.fac.ST_table.loc[(self.fac.ST_table["stage"] == self.stage)&\
                                               (self.fac.ST_table["from"]  == _from)&\
                                               (self.fac.ST_table["to"]    == to)]["setup_time"].values[0]
                    ST_list.append(ST)
            job_index = np.argmin(ST_list)
            job = self.space[job_index]
        # if self.fac.log:
        #     print('MST---')
        return job

    def EDD(self):
        job_index = np.argmin([job.DD for job in self.space])
        job = self.space[job_index]
        # if self.fac.log:
        #     print('EDD---')
        return job

    def SPT(self):
        job_index = np.argmin(job.PT[job.progress] for job in self.space)
        job = self.space[job_index]
        # if self.fac.log:
        #     print('SPT---')
        return job

    def dispatch(self, MC = None):
        if MC == None:
            for MC in self.OP:
                if len(self.space) <= 0:
                    break
                if MC.state == "idle":
                    if self.DP_rule == 'MST':
                        job = self.MST(MC)
                    elif self.DP_rule == 'EDD':
                        job = self.EDD()
                    elif self.DP_rule == 'SPT':
                        job = self.SPT()
                    else:
                        job = self.space[0]

                    if job != None:
                        MC.job_accept(job)
                        self.space.remove(job)
        else:
            if self.DP_rule == 'MST':
                job = self.MST(MC)
            elif self.DP_rule == 'EDD':
                job = self.EDD()
            elif self.DP_rule == 'SPT':
                job = self.SPT()
            else:
                job = self.space[0]

            if job != None:
                MC.job_accept(job)
                self.space.remove(job)

    def job_arrive(self, job):
        self.space.append(job)
        idle_machine = sum([1 if MC.state == "idle" else 0 for MC in self.OP])
        if idle_machine > 0:
            self.dispatch()

    def get_job(self, MC):
        if len(self.space) > 0:
            self.dispatch(MC)

class Machine:
    def __init__(self, fac, ID):#, mtype, stage):
        self.fac = fac
        # self.stage = stage
        # self.mtype = mtype
        self.ID = ID

    def initialize(self):#, IP, OP):
        #reference
        self.env = self.fac.env
        # self.IP = IP    #IP: input port, connect to a queue
        self.queues = self.fac.queues
        self.sink = self.fac.sink
        # self.OP = OP    #OP: output port, connect to a queue or sink
        #attribute
        self.space = []     #the job on machine put in space
        self.state = "idle"    # state is one of "idle", "setup", "processing"
        self.mode = None    #mode record the last processing job for checking the setup time
        #statistics
        self.SST = 0    #SST: start setup time
        self.SPT = 0    #SPT: start process time
        self.BT  = 0
        #initial process

    def job_accept(self, job):    #all job sould pass this operation not going stright to setup or process state
        self.space.append(job)
        #check setup requirement
        c1 = self.mode == None
        c1 = True
        if not c1:
            c2 = self.mode.jtype == job.jtype
        else:
            c2 = False
        if c1 or c2:    #don't need setup
                self.state = "processing"
                self.mode = self.space[0]
                self.SPT = self.env.now
                if self.fac.log:    #log message
                    print("{} : machine {} start processing job {}-{}".format(self.env.now, self.ID, self.space[0].ID, self.space[0].jtype))
                #set next event
                self.process = self.env.process(self.process_job())
        else:    #need setup
            self.state = "setup"
            self.SST = self.env.now
            if self.fac.log:    #log message
                print("{} : machine {} start setup {} -> {}".format(self.env.now, self.ID, self.mode.jtype, self.space[0].jtype))
            #set next evnet
            self.process = self.env.process(self.setup())

    def setup(self):
        #setnext event
        stage = 1#self.stage
        ID = self.ID
        _from = self.mode.jtype
        to = self.space[0].jtype
        setup_time = self.fac.ST_table[(self.fac.ST_table["stage"] == stage)&\
                                       (self.fac.ST_table["from"]  == _from)&\
                                       (self.fac.ST_table["to"]    == to)]["setup_time"].values[0]
        yield self.env.timeout(setup_time)
        #complete setup
        self.mode = self.space[0]
        '''
        setup stastic
        '''
        #save process information for gantt
        ST = self.SST
        FT = self.env.now
        MC = self
        job = "setup {} -> {}".format(_from, to)
        self.fac.update_gantt(MC, ST, FT, job)

        #process job
        self.state = "processing"
        self.SPT = self.env.now
        if self.fac.log:    #log message
            print("{} : machine {} start processing job {}-{}".format(self.env.now, self.ID, self.space[0].ID, self.space[0].jtype))
        #set next event
        self.process = self.env.process(self.process_job())
        '''
        event trace
        '''
        self.fac.print_state("{}_Setup_OK".format(self.ID))

    def process_job(self):
        job = self.space[0]
        #processing order for PT mins
        # PT = self.fac.PT_table.loc[(self.fac.PT_table["jtype"] == self.space[0].jtype)&\
        #                            (self.fac.PT_table["stage"] == self.stage)]["process_time"].values[0]
        print(len(job.PT), job.progress, job.ID, job.jtype)
        PT = job.PT[job.progress]
        yield self.env.timeout(PT)
        #change job progress
        job.progress += 1
        #complete process
        if self.fac.log:    #log message
            print("{} : machine {} finish processing job {}-{}".format(self.env.now, self.ID, self.space[0].ID, self.space[0].jtype))
        '''
        setup stastic
        '''
        #save process information for gantt
        ST = self.SPT
        FT = self.env.now
        MC = self
        job_now = "j{}-{}".format(self.space[0].ID, self.space[0].jtype)
        self.fac.update_gantt(MC, ST, FT, job_now)

        #for utilization
        self.BT += FT - ST

        #send job to next station
        if job.progress < len(job.routing):
            print(job.progress, len(job.routing))
            print(job.progress < len(job.routing))
            target = job.routing[job.progress]
            self.queues[target].job_arrive(job)
        else:
            self.sink.job_arrive(job)

        self.space = []
        #change state
        self.state = "idle"
        #get next job in queue
        self.queues[self.ID].get_job(self)
        '''
        event trace
        '''
        self.fac.print_state("{}_Complete".format(self.ID))

    def get_uti(self):
        busyTime = 0
        if self.state == "processing":
            busyTime = (self.env.now - self.SPT) + self.BT
        else:
            busyTime = self.BT
        uti = busyTime / self.env.now
        return uti

class Sink:
    def __init__(self, fac):
        self.fac = fac

    def initialize(self):
        #reference
        self.env = self.fac.env
        #attribute
        #statistics
        self.job_statistic = pd.DataFrame(columns = ["ID", "jtype", "arrival_time", "complete_time", "due_date", "flow_time", "tardiness", "lateness"])
        #initial process

    def job_arrive(self, job):
        #update factory statistic
        self.fac.throughput += 1
        self.fac.update_WIP(-1)
        #update job statistic
        job.CT = self.env.now
        self.update_job_statistic(job)
        # #ternimal condition
        # job_num = self.fac.job_info.shape[0]
        # if self.fac.throughput >= job_num:
        #     self.fac.terminal.succeed()

    def update_job_statistic(self, job):
        ID = job.ID
        AT = job.AT
        CT = job.CT
        DD = job.DD
        jtype = job.jtype
        flow_time = CT - AT
        lateness = CT - DD
        tardiness = max(0, lateness)
        self.job_statistic.loc[ID] = [ID, jtype, AT, CT, DD, flow_time, tardiness, lateness]

#factory
class Factory:
    def __init__(self, env, AV1, AV2, AV3, DD, PT_table, ST_table, DP_rule, log = True):
        self.env = env
        self.log = log
        self.AV1 = AV1
        self.AV2 = AV2
        self.AV3 = AV3
        self.DD_factor = DD
        self.PT_table = PT_table    #PT: process time
        self.ST_table = ST_table    #ST: setup time
        # self.ME_table = ME_table    #ME: machine eligibility
        self.DP_rule = DP_rule    #DP_rule: dispatching rule

        self.routing_dict = {"a": 'A,C,B,D',
                             "b": 'A,B,C',
                             "c": 'C,D,B'}
        self.PT_dict = {"a": '7,4,5,9',
                        "b": '6,6,6',
                        "c": '9,12,8'}
        ptA = sum(self.PT_table.loc[(self.PT_table["jtype"] == "A")]["process_time"])
        ptA = sum([int(i) for i in self.PT_dict["a"].split(',')])
        ptB = sum(self.PT_table.loc[(self.PT_table["jtype"] == "B")]["process_time"])
        ptB = sum([int(i) for i in self.PT_dict["b"].split(',')])
        ptC = sum(self.PT_table.loc[(self.PT_table["jtype"] == "C")]["process_time"])
        ptC = sum([int(i) for i in self.PT_dict["c"].split(',')])
        self.avg_pt  = {"a": ptA, "b": ptB, "c": ptC}

    def initialize(self):
        #build
        self.source = Source(self, self.AV1, self.AV2, self.AV3, self.DD_factor)
        self.queues = {"A" : Queue(self, 'A', self.DP_rule),
                       "B" : Queue(self, 'B', self.DP_rule),
                       "C" : Queue(self, 'C', self.DP_rule),
                       "D" : Queue(self, 'D', self.DP_rule)}
        self.machines = {"A" : Machine(self, 'A'),#, 1, 1),
                         "B" : Machine(self, 'B'),#, 1, 1),
                         "C" : Machine(self, 'C'),#, 1, 1),
                         "D" : Machine(self, 'D')}#, 1, 1),}
        self.sink = Sink(self)
        #initialize
        self.source.initialize()#self.queues["Q1"])

        for key, queue in self.queues.items():
            queue.initialize()
        for key, machine in self.machines.items():
            machine.initialize()

        self.sink.initialize()
        #attribute
        #statistics
        self.throughput = 0

        self.WIP = 0
        self.WIP_area = 0
        self.WIP_change_time = 0

        self.event_ID = 0
        self.event_record = pd.DataFrame(columns = ["type",
                                                    "time",
                                                    "queue A", "MC A",
                                                    "queue B", "MC B",
                                                    "queue C", "MC C",
                                                    "queue D", "MC D",
                                                    "WIP"])
        self.gantt_data = {"MC_name"         : [],
                           "start_process" : [],
                           "process_time"  : [],
                           "job"        : []}
        # #terminal event
        # self.terminal = self.env.event()

        self.print_state("Initializtion")

    def update_WIP(self, change):
        self.WIP_area += self.WIP*(self.env.now - self.WIP_change_time)
        self.WIP_change_time = self.env.now
        self.WIP += change

    def print_state(self, Etype):
        Q_A = "{}({})".format(len(self.queues["A"].space), [job.ID for job in self.queues["A"].space])
        Q_B = "{}({})".format(len(self.queues["B"].space), [job.ID for job in self.queues["B"].space])
        Q_C = "{}({})".format(len(self.queues["C"].space), [job.ID for job in self.queues["C"].space])
        Q_D = "{}({})".format(len(self.queues["D"].space), [job.ID for job in self.queues["D"].space])

        MC_A = self.machines["A"].state
        if MC_A != "idle":
            MC_A = "{}".format(self.machines["A"].space[0].ID)
        MC_B = self.machines["B"].state
        if MC_B != "idle":
            MC_B = "{}".format(self.machines["B"].space[0].ID)
        MC_C = self.machines["C"].state
        if MC_C != "idle":
            MC_C = "{}".format(self.machines["C"].space[0].ID)
        MC_D = self.machines["D"].state
        if MC_D != "idle":
            MC_D = "{}".format(self.machines["D"].space[0].ID)

        self.event_record.loc[self.event_ID] = [Etype,
                                                self.env.now,
                                                Q_A, MC_A,
                                                Q_B, MC_B,
                                                Q_C, MC_C,
                                                Q_D, MC_D,
                                                self.WIP]
        self.event_ID += 1

    def update_gantt(self, MC, ST, FT, job_now):
        self.gantt_data['MC_name'].append("M{}".format(MC.ID))
        self.gantt_data['start_process'].append(ST)
        self.gantt_data['process_time'].append(FT-ST)
        self.gantt_data['job'].append(job_now)

    def draw_gantt(self, save = None):
        fig, ax = plt.subplots(1,1)
        #set color list
        # colors = list(mcolors.CSS4_COLORS.keys())
        # np.random.shuffle(colors)
        colors = list(mcolors.TABLEAU_COLORS.keys())
        #draw gantt bar
        y = self.gantt_data['MC_name']
        width = self.gantt_data['process_time']
        left = self.gantt_data['start_process']
        def color_define(job, colors):
            if len(job) == 12:
                return 'black'
            else:
                if job[-1] == 'a':
                    return colors[1]
                elif job[-1] == 'b':
                    return colors[2]
                else:
                    return colors[3]

        color = [color_define(j, colors) for j in self.gantt_data['job']]
        ax.barh(y = y, width = width, height = 0.5, left = left, color = color, align = 'center', alpha = 0.6, edgecolor='black')
        #put the text on
        for i in range(len(self.gantt_data['MC_name'])):
            text_x = self.gantt_data['start_process'][i] + self.gantt_data['process_time'][i]/2
            text_y = self.gantt_data['MC_name'][i]
            text = self.gantt_data['job'][i]
            if len(text) >=10:
                text = ""
            ax.text(text_x, text_y, text, verticalalignment='center', horizontalalignment='center')
        #figure setting
        ax.set_xlabel("time (mins)")
        ax.set_xticks(np.arange(0, self.env.now+1, 20))
        ax.set_ylabel("Machine")
        ax.set_title("hybrid flow shop - gantt")

        if save is not None:
            fig.savefig(save)
        else:
            try:
                return fig
            except:
                plt.show()

    '''
    def print_schduling_data(self):
        df = pd.DataFrame(columns = ["job ID", "step", "MC ID", "start time", "end time"])
        for i in range(len(self.gantt_data["MC_id"])):
            order_id = self.gantt_data["order_id"][i]
            step = self.gantt_data["order_progress"][i]
            MC_id = self.gantt_data["MC_id"][i]
            ST = self.gantt_data["start_process"][i]
            ET = self.gantt_data["start_process"][i] + self.gantt_data["process_time"][i]
            df.loc[i] = [order_id, step, MC_id, ST, ET]
        with pd.option_context('display.max_rows', None, 'display.max_columns', df.shape[1]):
            print(df)
    '''

if __name__ == '__main__':
    #prepare the I/O path
    input_dir = os.getcwd() + '/input_data/'
    output_dir = os.getcwd() + '/results/'
    dir_list = [input_dir, output_dir]
    for dir in dir_list:
        if not os.path.exists(dir):
            os.makedirs(dir)
    #environment
    env = simpy.Environment()
    #parameter
    simulationTime = 100 #1440 * 5
    # job_info = pd.read_excel(input_dir + "job_information.xlsx")    # no used
    PT_table = pd.read_excel(input_dir + "PT_table.xlsx")   # should change to new
    ST_table = pd.read_excel(input_dir + "ST_table.xlsx")   # no used, change if need transition time
    # DP_rule = "FIFO"#"MST"

    # Decision case
    AV1 = int(input('>>> The average interarrival time of job type "A" = '))
    AV2 = int(input('>>> The average interarrival time of job type "B" = '))
    AV3 = int(input('>>> The average interarrival time of job type "C" = '))
    DD = float(input('>>> The due date tightness factor(k) = '))
    DP_rule = str(input('>>>[MST, EDD, SPT, FIFO] Choose one and key in!! >> '))

    #build job shop factory
    t0 = time.process_time()
    # print("t0 = {} seconds".format(t0))
    fac = Factory(env, AV1, AV2, AV3, DD, PT_table, ST_table, DP_rule, log = True)
    fac.initialize()
    #run simulation
    # env.run(until = fac.terminal)
    env.run(until = simulationTime)

    fac.source.df.to_excel("job_info.xlsx")
    #print result
    print("====================================================================")
    print("dispatching rule: {}".format(DP_rule))
    print(fac.event_record)
    print("==================================================================================")
    # with pd.option_context('display.max_rows', None, 'display.max_columns', None):#fac.sink.job_statistic.shape[1]):
    print(fac.sink.job_statistic.sort_values(by = "ID"))

    print("==================================================================================")
    print("throughput: {}".format(fac.throughput))
    print("makespan: {}".format(fac.env.now))
    print("Average flow time: {}".format(np.mean(fac.sink.job_statistic['flow_time'])))
    print("Average tardiness: {}".format(np.mean(fac.sink.job_statistic['tardiness'])))
    print("Average lateness: {}".format(np.mean(fac.sink.job_statistic['lateness'])))
    print("Aveage WIP: {}".format(fac.WIP_area/fac.env.now))
    print("====================================================================")
    uti_lst = []
    for key, mc in fac.machines.items():
        uti = mc.get_uti()
        uti_lst.append(uti)
        print("Utilization of {}: {}".format(key, uti))
    uti_lst = np.array(uti_lst)
    print("Average utilization: {}".format(np.mean(uti_lst)))
    print("====================================================================")

    cpu_time = time.process_time() - t0
    print("Cost {} seconds".format(cpu_time))

    '''
    print("schduling data")
    fac.print_schduling_data()
    print("====================================================================")
    '''
    fac.draw_gantt(output_dir + "{}.png".format(DP_rule))
    # plt.show()


from typing import Tuple

import pytest
from ground.hints import Scalar
from hypothesis import given
from hypothesis.errors import HypothesisWarning
from hypothesis.strategies import DataObject

from hypothesis_geometry.hints import Strategy
from hypothesis_geometry.planar import contours
from tests import strategies
from tests.utils import (ScalarsLimitsType,
                         SizesPair,
                         are_vertices_strict,
                         contour_has_coordinates_in_range,
                         contour_has_coordinates_types,
                         contour_has_valid_sizes,
                         is_contour,
                         is_contour_counterclockwise,
                         is_contour_non_self_intersecting,
                         is_contour_strict)

@given(strategies.scalars_strategies,
       strategies.convex_contours_sizes_pairs)
def test_basic(coordinates: Strategy[Scalar],
               sizes_pair: SizesPair) -> None:
    min_size, max_size = sizes_pair

    result = contours(coordinates,
                      min_size=min_size,
                      max_size=max_size)

    assert isinstance(result, Strategy)

@given(strategies.data,
       strategies.scalars_strategy_with_limit_and_type_pairs,
       strategies.convex_contours_sizes_pairs)
def test_properties(data: DataObject,
                    coordinates_limits_type_pair: Tuple[ScalarsLimitsType,
                                                        ScalarsLimitsType],
                    sizes_pair: SizesPair) -> None:
    (x_coordinates_limits_type,
     y_coordinates_limits_type) = coordinates_limits_type_pair
    ((x_coordinates, (min_x_value, max_x_value)),
     x_type) = x_coordinates_limits_type
    ((y_coordinates, (min_y_value, max_y_value)),
     y_type) = y_coordinates_limits_type
    min_size, max_size = sizes_pair

    strategy = contours(x_coordinates, y_coordinates,
                        min_size=min_size,
                        max_size=max_size)

    result = data.draw(strategy)

    assert is_contour(result)
    assert contour_has_valid_sizes(result,
                                   min_size=min_size,
                                   max_size=max_size)
    assert contour_has_coordinates_types(result,
                                         x_type=x_type,
                                         y_type=y_type)
    assert contour_has_coordinates_in_range(result,
                                            min_x_value=min_x_value,
                                            max_x_value=max_x_value,
                                            min_y_value=min_y_value,
                                            max_y_value=max_y_value)
    assert is_contour_strict(result)
    assert is_contour_non_self_intersecting(result)
    assert is_contour_counterclockwise(result)

@given(strategies.data,
       strategies.scalars_strategies_with_limits_and_types,
       strategies.convex_contours_sizes_pairs)
def test_same_coordinates(data: DataObject,
                          coordinates_limits_type: ScalarsLimitsType,
                          sizes_pair: SizesPair) -> None:
    (coordinates, (min_value, max_value)), type_ = coordinates_limits_type
    min_size, max_size = sizes_pair

    strategy = contours(coordinates,
                        min_size=min_size,
                        max_size=max_size)

    result = data.draw(strategy)

    assert is_contour(result)
    assert contour_has_valid_sizes(result,
                                   min_size=min_size,
                                   max_size=max_size)
    assert contour_has_coordinates_types(result,
                                         x_type=type_,
                                         y_type=type_)
    assert contour_has_coordinates_in_range(result,
                                            min_x_value=min_value,
                                            max_x_value=max_value,
                                            min_y_value=min_value,
                                            max_y_value=max_value)
    assert are_vertices_strict(result.vertices)
    assert is_contour_non_self_intersecting(result)
    assert is_contour_counterclockwise(result)

@given(strategies.scalars_strategies,
       strategies.invalid_convex_contours_sizes_pairs)
def test_invalid_sizes(coordinates: Strategy[Scalar],
                       invalid_sizes_pair: SizesPair) -> None:
    min_size, max_size = invalid_sizes_pair

    with pytest.raises(ValueError):
        contours(coordinates,
                 min_size=min_size,
                 max_size=max_size)

@given(strategies.scalars_strategies,
       strategies.non_valid_convex_contours_sizes_pairs)
def test_non_valid_sizes(coordinates: Strategy[Scalar],
                         non_valid_sizes_pair: SizesPair) -> None:
    min_size, max_size = non_valid_sizes_pair

    with pytest.warns(HypothesisWarning) as warnings:
        contours(coordinates,
                 min_size=min_size,
                 max_size=max_size)

    assert len(warnings) == 1


from detectron2.structures import BoxMode
from detectron2.data import build_detection_train_loader, build_detection_test_loader
from detectron2.data import transforms as T
from detectron2.data import detection_utils as utils
from detectron2.engine import DefaultTrainer
from detectron2.utils.visualizer import Visualizer, ColorMode

import torch
import cv2

import copy
import os
import xml.etree.ElementTree as elemTree

import math
import numpy as np
from pytesseract import image_to_string
from PIL import Image
import pandas
import csv

text_dir = '/home/ubuntu/data/Workspace/Soobin/wide_report/classes.txt'
# category 개수 count
def category_cnt(text_dir=text_dir):
    f = open(text_dir, "r")
    category_names = f.readlines()
    return len(category_names)
classes=category_cnt()

# xml file (annotation) 을 불러와 detectron2의 모델에 맞게 data 입력
def get_data_dicts(img_dir, annotation_dir, text_dir):
    dataset_dicts = []
    file_list = os.listdir(img_dir)
    idx = 0

    f = open(text_dir, "r")
    category_names = f.readlines()
    category_dict = {}
    # del category_names[0]
    category_idx = 0
    for category_name in category_names:
        category_dict[category_name.split("\n")[0]] = category_idx
        category_idx += 1

    for file_name in file_list:
        record = {}

        tree = elemTree.parse(os.path.join(annotation_dir, file_name.split(".")[0] + ".xml"))
        root = tree.getroot()

        width = int(root[4][0].text)
        height = int(root[4][1].text)
        idx += 1

        record["file_name"] = img_dir + file_name
        record["image_id"] = idx
        record["width"] = width
        record["height"] = height

        objects = []
        for i in range(len(root.findall("object"))):
            category = root.findall("object")[i].find("name").text
            category_idx = category_dict[category]
            xmin = int(root.findall("object")[i].find("bndbox").find("xmin").text)
            ymin = int(root.findall("object")[i].find("bndbox").find("ymin").text)
            xmax = int(root.findall("object")[i].find("bndbox").find("xmax").text)
            ymax = int(root.findall("object")[i].find("bndbox").find("ymax").text)

            obj = {
                "bbox":[xmin, ymin, xmax, ymax],
                "bbox_mode":BoxMode.XYXY_ABS,
                "category_id":category_idx,
                "iscrowd":0
            }
            objects.append(obj)
        record["annotations"] = objects
        dataset_dicts.append(record)

    return dataset_dicts

def get_predict_dicts(img_dir):
    dataset_dicts = []
    file_list = os.listdir(img_dir)
    idx = 0

    for file_name in file_list:
        record = {}

        height, width = cv2.imread(img_dir + file_name).shape[:2]
        idx += 1

        record["file_name"] = img_dir + file_name
        record["image_id"] = idx

        dataset_dicts.append(record)

    return dataset_dicts

# data annotation이 필요한 경우 transformation 함수를 사용하기 위해 설정해 놓음
def my_mapper(dataset_dict):
    dataset_dict = copy.deepcopy(dataset_dict)
    image = utils.read_image(dataset_dict["file_name"], format="BGR")
    image, transforms = T.apply_transform_gens([
        T.RandomExtent([0.8, 1.2], [1.5, 1.5]),
    ], image)
    dataset_dict["image"] = torch.as_tensor(image.transpose(2,0,1).astype("float32"))

    annos = [
        utils.transform_instance_annotations(obj, transforms, image.shape[:2])
        for obj in dataset_dict.pop("annotations")
        if obj.get("iscrowd", 0) == 0
    ]
    instances = utils.annotations_to_instances(annos, image.shape[:2])
    dataset_dict["instances"] = utils.filter_empty_instances(instances)
    return dataset_dict

class My_Trainer(DefaultTrainer):
    @classmethod
    def build_test_loader(cls, cfg, dataset_name):
        return build_detection_test_loader(cfg, dataset_name, mapper=None)

    @classmethod
    def build_train_loader(cls, cfg):
        return build_detection_train_loader(cfg, mapper=my_mapper)

# get grayscale image
def get_grayscale(image):
    return cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

# noise removal
def remove_noise(image):
    return cv2.medianBlur(image,5)

#thresholding
def thresholding(image):
    return cv2.threshold(image, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)[1]

#erosion
def erode(image):
    kernel = np.ones((3,3),np.uint8)
    return cv2.erode(image, kernel, iterations = 1)

#opening - erosion followed by dilation
def opening(image):
    kernel = np.ones((5,5),np.uint8)
    return cv2.morphologyEx(image, cv2.MORPH_OPEN, kernel)

#canny edge detection
def canny(image):
    return cv2.Canny(image, 100, 200)

def image_processing(img):
    gray = get_grayscale(img)
    # thresh = canny(gray)
    # clean = remove_noise(gray)
    # thred = thresholding(gray)
    # eroding = erode(gray)
    return gray

# 각 box들을 추출 후 저장
def get_output_image(predictor, image_dir, metadata, text_dir, expnum):
    image_list = os.listdir(image_dir)

    f = open(text_dir, "r")
    category_names = f.readlines()
    category_dict = {}
    # del category_names[0]
    category_idx = 0
    for category_name in category_names:
        category_dict[category_idx] = category_name.split("\n")[0]

        if not os.path.exists("./img%02i/category%02i_%s/"%(expnum, category_idx,category_dict[category_idx])):
            os.makedirs("./img%02i/category%02i_%s/"%(expnum, category_idx,category_dict[category_idx]))

        category_idx += 1

    for image_name in image_list:
        image = cv2.imread(image_dir + image_name)
        outputs = predictor(image)

        v = Visualizer(image[:,:,::-1],
                    metadata=metadata,
                    scale=1.0,
                    instance_mode=ColorMode.IMAGE)
        v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
        cv2.imwrite("./img%02i/%s"%(expnum, image_name), v.get_image()[:,:,::-1])

        boxes = outputs["instances"].pred_boxes.tensor.cpu().numpy()
        classes = outputs["instances"].pred_classes.cpu().numpy()

        for i in range(len(boxes)):
            pred_box = boxes[i]
            pred_class = classes[i]
            # if category_dict[pred_class]=='erase':
            #     erase = image.copy()
            #     erase[int(pred_box[1]):int(pred_box[3]), int(pred_box[0]):int(pred_box[2])] = (255,255,255)
            #     cv2.imwrite("./img%02i/category%02i_%s/%s"%(expnum, pred_class, category_dict[pred_class], image_name), erase)
            #     continue
            cropped_img = image[int(pred_box[1]):int(pred_box[3]), int(pred_box[0]):int(pred_box[2])]
            saving_img = image_processing(cropped_img)
            cv2.imwrite("./img%02i/category%02i_%s/%s"%(expnum, pred_class, category_dict[pred_class], image_name), saving_img)

def rename_file(predictor, image_dir, metadata, text_dir):
    image_list = os.listdir(image_dir)
    f = open(text_dir, "r")
    category_names = f.readlines()
    category_dict = {}
    # del category_names[0]
    category_idx = 0
    for category_name in category_names:
        category_dict[category_idx] = category_name.split("\n")[0]
        category_idx += 1

    for image_name in image_list:
        image = cv2.imread(image_dir + image_name)
        outputs = predictor(image)

        v = Visualizer(image[:,:,::-1],
                    metadata=metadata,
                    scale=1.0,
                    instance_mode=ColorMode.IMAGE)
        v = v.draw_instance_predictions(outputs["instances"].to("cpu"))

        boxes = outputs["instances"].pred_boxes.tensor.cpu().numpy()
        classes = outputs["instances"].pred_classes.cpu().numpy()

        for i in range(len(boxes)):
            pred_box = boxes[i]
            pred_class = classes[i]
            if category_dict[pred_class]=='side':
                cropped = image[int(pred_box[1]):int(pred_box[3]), int(pred_box[0]):int(pred_box[2])]
                text = image_to_string(cropped,config="-c tessedit_char_whitelist=(ODSLR) -psm 6")
                print(text)

                if 'R' in text:
                    os.rename(image_dir+image_name,image_dir+image_name[:-4]+'_dl_OD.jpg')
                    print(image_dir+image_name,image_dir+image_name[:-4]+'_dl_OD.jpg')
                elif 'L' in text:
                    os.rename(image_dir+image_name,image_dir+image_name[:-4]+'_dl_OS.jpg')
                    print(image_dir+image_name,image_dir+image_name[:-4]+'_dl_OS.jpg')
                else:
                    continue

def del_custom_id_rename(predictor, image_dir, metadata, expnum, eye_dict):
    image_list = os.listdir(image_dir)

    if not os.path.exists("./img%02i/del_custom_id/"%(expnum)):
        os.mkdir("./img%02i/del_custom_id/"%(expnum))

    for image_name in image_list:
        image = cv2.imread(image_dir + image_name)
        outputs = predictor(image)

        v = Visualizer(image[:,:,::-1],
                    metadata=metadata,
                    scale=1.0,
                    instance_mode=ColorMode.IMAGE)
        v = v.draw_instance_predictions(outputs["instances"].to("cpu"))

        boxes = outputs["instances"].pred_boxes.tensor.cpu().numpy()
        classes = outputs["instances"].pred_classes.cpu().numpy()

        eye = eye_dict[image_name]

        if eye == 0:
            eye = "OD"  # 오른쪽
        elif eye == 1:
            eye = "OS"  # 왼  쪽

        box_dict = {}
        for index in range(len(boxes)):
            box_dict[classes[index]] = boxes[index]

        for i in [0, 1, 3, 15, 25]:
            if i in box_dict.keys():
                pred_box = box_dict[i]

                image[int(pred_box[1]):int(pred_box[3]), int(pred_box[0]):int(pred_box[2])] = image[1,1]
                cv2.imwrite("./img%02i/del_custom_id/%s_dl_%s.jpg"%(expnum, image_name.split(".")[0],eye), image)

def read_OCT_master_table():
    table_dir='/home/ubuntu/data/Workspace/Soobin/data/frequent_list/OCT var master tableV2.csv'
    master_table = pandas.read_csv(table_dir)
    return master_table

def num_extract(path,img):
    width = 60
    height=30
    top = int(height*0.15)
    bottom=top
    left = int(width*0.15)
    right=left
    dim = (width,height)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    img = cv2.resize(img,dim,interpolation=cv2.INTER_NEAREST)
    blur = cv2.GaussianBlur(img,(3,3),0)
    ret, temp = cv2.threshold(blur, int(cv2.mean(blur)[0]), 255,cv2.THRESH_TRUNC)
    ret, thr2 = cv2.threshold(temp, int(cv2.mean(temp)[0]), 255, cv2.THRESH_BINARY)
    im2 = cv2.copyMakeBorder(thr2,top,bottom,left,right,cv2.BORDER_CONSTANT,None,[255,255,255])
    # cv2.imwrite(path+'_processed.jpg',im2)
    text = image_to_string(im2, config="-c tessedit_char_whitelist=0123456789 -psm 6")
    return text

master_table = read_OCT_master_table()
def add_csv():
    mat = master_table

def extract_txt_add_csv(image_dir, text_dir, expnum):
    master_table = read_OCT_master_table()

    image_list = os.listdir(image_dir)
    f = open(text_dir, "r")
    category_names = f.readlines()
    category_dict = {}
    category_list=[]

    category_idx = 0
    for category_name in category_names:
        category_dict[category_idx] = category_name.split("\n")[0]
        category_idx += 1

    for category_idx in range(0,classes):
        category_list.append(category_dict[category_idx])

    for image_name in image_list:
        for category_idx in range(0,classes):
            image_file = "./img%02i/category%02i_%s/%s"%(expnum, category_idx, category_dict[category_idx], image_name)
            if os.path.exists(image_file):
                txt_file = open("./img%02i/category%02i_%s/%s.txt"%(expnum, category_idx, category_dict[category_idx], image_name),'w')
                image = cv2.imread(image_file)
                if category_idx in [4,5,6,7,8,9,10,11,12]:
                    # text = num_preprocess_extract(image_file)        # dpi upgraded
                    text = num_extract(image_file,image)
                elif category_idx in [13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,37]:
                    text = image_to_string(image, config="-c tessedit_char_whitelist=0123456789 -psm 6")
                else:
                    text = image_to_string(image)

                    # if 'R' is in text:
                    #     rename_official
                txt_file.write(text)
                txt_file.close()

def rename_official(image_dir,expnum):
    for image_name in image_list:
        csv_save_list = []
        txt_file = open("./img%02i/text/%s.txt"%(expnum, image_name.split(".")[0]), 'w')
        csv_save_list.append(image_name.split(".")[0])
        for category_idx in range(0,classes):
            image_file = "./img%02i/category%02i_%s/%s"%(expnum, category_idx, category_dict[category_idx], image_name)
            if os.path.exists(image_file):
                image = cv2.imread(image_file)
                resize_image = cv2.resize(image, dsize=(0,0), fx=4.3, fy=4.3, interpolation=cv2.INTER_LANCZOS4)
                text = image_to_string(resize_image)
                # csv_save_list.append(text)
                txt_file.write(text)

            else:
                continue

            if category_idx == 2:
                if "R" in text:
                    eye_dict[image_name] = 0
                else:
                    eye_dict[image_name] = 1

# 추출한 text data를 .txt file로 저장
def generate_txt_data(image_dir, text_dir, expnum, times):
    f = open(text_dir, "r")
    category_names = f.readlines()
    category_list=["file_name"]
    category_dict = {}
    # del category_names[0]
    category_idx = 0
    for category_name in category_names:
        category_dict[category_idx] = category_name.split("\n")[0]
        category_idx += 1

    for category_idx in range(0,classes):
        category_list.append(category_dict[category_idx])

    if not os.path.exists("./img%02i/text/"%(expnum)):
        os.mkdir("./img%02i/text/"%(expnum))

    image_list = os.listdir(image_dir)

    eye_dict = {}

    csv_total_list =[]

    for image_name in image_list:
        csv_save_list = []
        txt_file = open("./img%02i/text/%s.txt"%(expnum, image_name.split(".")[0]), 'w')
        csv_save_list.append(image_name.split(".")[0])
        for category_idx in range(0,classes):
            image_file = "./img%02i/category%02i_%s/%s"%(expnum, category_idx, category_dict[category_idx], image_name)
            if os.path.exists(image_file):
                image = cv2.imread(image_file)
                resize_image = cv2.resize(image, dsize=(0,0), fx=4.3, fy=4.3, interpolation=cv2.INTER_LANCZOS4)
                text = image_to_string(resize_image)
                # csv_save_list.append(text)
                txt_file.write(text)

            else:
                continue

            if category_idx == 2:
                if "R" in text:
                    eye_dict[image_name] = 0
                else:
                    eye_dict[image_name] = 1

        csv_total_list.append(csv_save_list)

    # csv_file = pandas.DataFrame(csv_total_list, columns=category_list)
    # csv_file.to_csv("./img%02i/VF_text_result.csv"%(expnum), index=False, encoding='utf-8')

        # txt_file.close()

    return eye_dict


#!/usr/bin/python

""" Pourbaix plot """

import numpy as np
from useful_functions import AutoVivification
from pprint import pprint
from ase.db import connect
import os, csv, sys
from parser_class import ParseInfo
import matplotlib
import matplotlib.image as mpimg
import matplotlib.pyplot as plt
from plot_params import get_plot_params_Times

def main(thermodb, referdb, list_cells, facet, functional):
    pour = ParseInfo(thermodb, referdb, list_cells, facet, functional,ref_type='Pb')
    pour.get_pourbaix()
    return pour.min_energy_cell

if __name__ == '__main__':
    ## Details
    output = 'output/'
    os.system('mkdir -p ' + output)

    get_plot_params_Times()

    # facets
    facets = ['111', '100', '110',  '211', ] # 'recon_110'
    functional = 'BF'
    no_electrons_Pb = 2 # Two electrons are transferred
    U_PB_2_SHE = -0.13 # V To convert to SHE

    # databases
    thermodb = connect('databases/Au_Pb_coverage.db')
    referdb = connect('databases/reference_lead.db')

    # Stylistic
    colors_coverage = ['tab:blue', 'tab:green', 'tab:red']
    cell_sizes = {
                  '100': ['1x1', '2x2', '3x3'],
                  '211': ['1x3', '2x3', '3x3'],
                  '111': ['1x1', '2x2', '3x3'],
                  '110': ['1x1', '2x2', '3x3'],
                  'recon_110': ['1x1', '2x1', '3x1'],
                  }

    j_ylim = {'100':[-70, 70], '111':[-120, 120], '211':[-70, 70], \
              '110':[-70, 70], 'recon_110':[-70, 70],}

    cell_mult = {
                  '100': {'1x1':1, '2x2':1/4, '3x3':1/9},
                  '211': {'1x3':1, '2x3':1/2, '3x3':1/3},
                  '111': {'1x1':1, '2x2':1/4, '3x3':1/9},
                  '110': {'1x1':1, '2x2':1/4, '3x3':1/9},
                  'recon_110':{'1x1':1, '2x1':1/2, '3x1':1/3},
                }

    coverage_labels = {
                  '100': {'1x1':'1', '2x2':r'$\frac{1}{4}$', '3x3':r'$\frac{1}{9}$'},
                  '211': {'1x3':'1', '2x3':r'$\frac{2}{3}$', '3x3':r'$\frac{1}{3}$'},
                  '111': {'1x1':'1', '2x2':r'$\frac{1}{4}$', '3x3':r'$\frac{1}{9}$'},
                  '110': {'1x1':'1', '2x2':r'$\frac{1}{4}$', '3x3':r'$\frac{1}{9}$'},
                  'recon_110':{'1x1':'1', '2x1':r'$\frac{1}{2}$', '3x1':r'$\frac{1}{3}$'},
                  }

    coverages_cell = {
                      '100': [1, 0.25, 0.11],
                      '211': [1, 0.66, 0.33],
                      '111': [1, 0.25, 0.11],
                      '110': [1, 0.25, 0.11],
                      'recon_110':[1, 0.25, 0.33]
                      }

    alphabets = ['a', 'b', 'c', 'd', 'e', 'f', 'g']

    facet_markers = {'100':'o', '111':'o', '211':'o', '110':'o'}
    potential_range = np.linspace(-1.4, 1.4)

    # Plot related
    fig, ax1 = plt.subplots(len(facets)+2, 2, sharex=True, figsize=(12, 24))
    gs = ax1[-1, -1].get_gridspec()
    for ax in ax1[len(facets)+1, : ]:
        ax.remove()
    for ax in ax1[len(facets),:]:
        ax.remove()

    # Add images of different coverages to plot
    axbig111 = fig.add_subplot(gs[len(facets),:])
    axbig211 = fig.add_subplot(gs[len(facets)+1,:])

    img_211 = mpimg.imread('blender_images/Au_Pb_211.png')
    axbig211.imshow(img_211)

    img_111 = mpimg.imread('blender_images/Au_Pb_111.png')
    axbig111.imshow(img_111)

    # add annotations
    axbig211.annotate(r'$\theta = 1ML $', xy=(0.125, 0.1), xycoords="axes fraction", color=colors_coverage[0])
    axbig211.annotate(r'$\theta = \frac{2}{3}ML $', xy=(0.45, 0.1), xycoords="axes fraction", color=colors_coverage[1])
    axbig211.annotate(r'$\theta = \frac{1}{3}ML $', xy=(0.80, 0.1), xycoords="axes fraction", color=colors_coverage[2])
    axbig211.axis('off')
    axbig211.set_title(r'Au Steps')
    axbig211.annotate('f)', xy=(-0.1, 1.1),xycoords="axes fraction", fontsize=32)

    axbig111.annotate(r'$\theta = 1ML $', xy=(0.05, -0.1), xycoords="axes fraction", color=colors_coverage[0])
    axbig111.annotate(r'$\theta = \frac{1}{4}ML $', xy=(0.37, -0.1), xycoords="axes fraction", color=colors_coverage[1])
    axbig111.annotate(r'$\theta = \frac{1}{9}ML $', xy=(0.65, -0.1), xycoords="axes fraction", color=colors_coverage[2])
    axbig111.axis('off')
    axbig111.set_title(r'Au Terraces')
    axbig111.annotate('e)', xy=(-0.1, 1.1),xycoords="axes fraction", fontsize=32)
    # Store data
    results = AutoVivification()
    experiments = AutoVivification()

    ################################################
    # STORE DATA

    for facet in facets:
        results[facet] = main(thermodb, referdb, cell_sizes[facet], facet, functional)
        data_facet = []
        with open('previous/Au' + facet + '.csv') as f:
            #experiments[facet] = np.loadtxt('previous/Au' + facet + '.csv', delimiter=',')
            read_f = csv.reader(f,delimiter=';', dialect=csv.excel,)
            retry = False
            for row in read_f:
                try:
                    data_facet.append([float(a) for a in row])
                except ValueError:
                    retry = True
            # if retry:
            #     read_f = csv.reader(f,delimiter=',', dialect=csv.excel,)
            #     for row in read_f:
            #         data_facet.append([float(a) for a in row])

        experiments[facet] = np.array(data_facet).transpose()

    ################################################
    # PLOT DATA

    # fig.subplots_adjust(hspace=0)
    for ind, facet in enumerate(facets):
        #plt.figure()
        p_all = []
        for index, cell in enumerate(results[facet]):
            nPb = cell_mult[facet][cell]
            # Straight line slope for Pourbaix diagram
            p = np.poly1d([2*nPb, nPb * results[facet][cell]  - 2 * nPb * U_PB_2_SHE])
            if index == len(results[facet])-1:
                p_all.append(np.poly1d([0]))
            p_all.append(p)

            # Plot the experimental plot
            ax1[ind,1].plot(experiments[facet][0], experiments[facet][1], color='tab:gray', alpha=0.5,lw=4)
            ax1[ind,0].plot(potential_range, p(potential_range)  ,
                    color=colors_coverage[index], lw=4,  label=r'$\theta = $ ' + coverage_labels[facet][cell] + ' ML')

            # ax1[0,ind].tick_params(axis='both', which='major', labelsize=22)
            # if ind == 0:
            ax1[ind,0].set_ylabel(r'$\Delta E_{Pb}$ / eV', fontsize=32)
            ax1[ind,1].set_ylabel(r'$j$ / $\mu A cm^{-2}$', fontsize=32)
            if ind == len(facets)-1:
                ax1[ind,1].set_xlabel(r'Potential vs SHE / V', fontsize=32)
                ax1[ind,0].set_xlabel(r'Potential vs SHE / V', fontsize=32)
            ax1[ind,0].set_ylim([-0.7, 0.7])
            ax1[ind,0].set_xlim([-0.7, 0.7])
            # ax1[0,ind].legend(bbox_to_anchor=(0,1.02,1,0.2), loc="lower left",
                # mode="expand", borderaxespad=0, ncol=3, fontsize=22)
            if ind == 0:
                ax1[ind,0].annotate(r'$\mathbf{\theta} = $ ' + coverage_labels[facet][cell] + ' ML',
                                xy=(0.65, 0.05+  0.15 * index), backgroundcolor="w", \
                                    xycoords="axes fraction", color=colors_coverage[index]).draggable()
            else:
                ax1[ind,0].annotate(r'$\mathbf{\theta} = $ ' + coverage_labels[facet][cell] + ' ML',
                                xy=(0.03, 0.5 + 0.15 * index), backgroundcolor="w", \
                                     xycoords="axes fraction", color=colors_coverage[index]).draggable()
            ax1[ind,0].xaxis.set_ticks_position('bottom')
            ax1[ind,1].xaxis.set_ticks_position('bottom')

            ax1[ind,1].annotate(r'Au(' +facet.replace('recon_', 'recons-') + ')', xy=(0.03, 0.5), \
                    xycoords="axes fraction", fontsize=28, color='tab:brown', weight='bold')

            ax1[ind,0].annotate(alphabets[ind] + ')', xy=(-0.1, 1.1),xycoords="axes fraction", fontsize=32)

            # ax1[ind,0].tick_params(axis='both', which='major', labelsize=22)
            # ax1[ind,1].tick_params(axis='both', which='major', labelsize=22)
            ax1[ind,1].set_ylim(j_ylim[facet])

        ax1[ind,0].axhline(y=0, color='k', ls='-', lw=4, )

        for i in range(len(p_all)-1):
            points_inter = (p_all[i+1] - p_all[i]).r
            ax1[ind,0].axvline(x=points_inter, ls='--', color='grey')
            ax1[ind,1].axvline(x=points_inter, ls='--', color='grey')

    plt.tight_layout()
    plt.savefig(output + 'lead_UPD.pdf', )
    plt.savefig(output + 'lead_UPD.png', )
    plt.show()


#!/usr/bin/env python
'''
Created on Nov 14, 2012

@author: shun
'''

import logging
import sys
import os

LOG_FORMAT = '%(threadName)s %(levelname)-5s %(asctime)s %(name)s %(message)s'

def config_logging(filepath):
    logging.basicConfig(filename=filepath, level=logging.DEBUG, format=LOG_FORMAT)

def config_as_default():
    config_logging(os.devnull)

def setLevel(namespace, level):
    logger = logging.getLogger(namespace)
    logger.setLevel(getattr(logging, level))

def getLogger(name, filename=None):
    if not filename:
        filename = name

    LOG = logging.getLogger(name)
    formatter = logging.Formatter(LOG_FORMAT)

    def _stdout_handler():
        infohand = logging.StreamHandler(sys.stdout)
        infohand.setLevel(logging.INFO)
        infohand.setFormatter(formatter)
        return infohand

    def _file_handler():
        filehand = logging.FileHandler(filename, 'a')
        filehand.setLevel(logging.DEBUG)
        filehand.setFormatter(formatter)
        return filehand

    LOG.addHandler(_stdout_handler())
    LOG.addHandler(_file_handler())
    return LOG

if __name__ == '__main__':
    logger = getLogger('fuck', 'fuck.log')
    logger.info('info...')
    logger.debug('debug...')


import inject
import json
import pytest

from store_app import app
from fake_storage import FakeStorage

def configure_test(binder):
    db = FakeStorage()
    binder.bind('DB', db)

class Initializer:
    def setup(self):
        inject.clear_and_configure(configure_test)
        app.config['TESTING'] = True
        with app.test_client() as client:
            self.client = client

class TestUsers(Initializer):
    def test_create_new(self):
        resp = self.client.post(
            '/users',
            json={'name': 'John Doe'}
        )
        assert resp.status_code == 201
        assert resp.json == {'user_id': 1}
        resp = self.client.post(
            '/users',
            json={'name': 'Andrew Derkach'}
        )
        assert resp.status_code == 201
        assert resp.json == {'user_id': 2}

    def test_successful_get_user(self):
        resp = self.client.post(
            '/users',
            json={'name': 'John Doe'}
        )
        user_id = resp.json['user_id']
        resp = self.client.get(f'/users/{user_id}')
        assert resp.status_code == 200
        assert resp.json == {'name': 'John Doe'}

    def test_get_unexistant_user(self):
        resp = self.client.get(f'/users/1')
        assert resp.status_code == 404
        assert resp.json == {'error': 'No such user_id 1'}

    def test_succesful_put_user(self):
        resp = self.client.post(
            '/users',
            json={'name': 'John Doe'}
        )
        user_id = resp.json['user_id']
        resp = self.client.put(
            f'/users/{user_id}',
            json={'name': 'Johanna Doe'}
        )
        assert resp.status_code == 200
        assert resp.json == {'status': 'success'}

    def test_unexistant_put_user(self):
        resp = self.client.put(
            f'/users/56',
            json={'name': 'John Doe'}
        )
        assert resp.status_code == 404

class TestGoods(Initializer):
    def test_post_goods(self):
        with open('jsons/goods.json') as goods:
            resp = self.client.post(
                '/goods',
                json=json.load(goods)
            )
        assert resp.status_code == 201
        assert resp.json == {'numbers of items created': 10}

    def test_get_goods(self):
        with open('jsons/goods.json') as goods:
            self.client.post(
                '/goods',
                json=json.load(goods)
            )
        resp = self.client.get('/goods')
        assert resp.status_code == 200
        with open('jsons/goods_updated.json') as new_goods:
            assert resp.json == json.load(new_goods)

    def test_update_goods(self):
        with open('jsons/goods.json') as goods:
            self.client.post(
                '/goods',
                json=json.load(goods)
            )
        with open('jsons/goods_for_update.json') as update:
            resp = self.client.put(
                '/goods',
                json=json.load(update)
            )
            assert resp.json == {'successfully_updated': 3,
                                 'errors': {'no such id in goods': [11, 12, 13]}}
        resp = self.client.get('/goods')
        assert resp.status_code == 200
        with open('jsons/goods_updated_2.json') as new_goods:
            assert resp.json == json.load(new_goods)

class TestStores(Initializer):
    @pytest.mark.parametrize(
        "store,answer_code,answer",
        (
                ({'name': 'Mad Cow', 'location': 'Lviv', 'manager_id': 1}, 201, {'store_id': 1}),
                ({'name': 'Mad Cow', 'location': 'Lviv', 'manager_id': 15}, 404, {'error': 'No such user_id 15'})
        )
    )
    def test_post_store(self, store, answer_code, answer):
        user = {'name': 'John Doe'}
        self.client.post('/users', json=user)
        resp = self.client.post('/stores', json=store)
        assert resp.status_code == answer_code
        assert resp.json == answer

    @pytest.mark.parametrize(
        "adress,answer_code,answer",
        (
                (1, 200, {'name': 'Mad Cow', 'location': 'Lviv', 'manager_id': 1}),
                (2, 404, {'error': 'No such store_id 2'})
        )
    )
    def test_get_store_by_id(self, adress, answer_code, answer):
        user = {'name': 'John Doe'}
        self.client.post('/users', json=user)
        newstore = {'name': 'Mad Cow', 'location': 'Lviv', 'manager_id': 1}
        self.client.post('/stores', json=newstore)
        resp = self.client.get(f'/stores/{adress}')
        assert resp.status_code == answer_code
        assert resp.json == answer

    @pytest.mark.parametrize(
        "update,adress,answer_code,answer",
        (
                ({'name': 'Local Taste', 'location': 'Lviv', 'manager_id': 1}, 1, 201, {'status': 'success'}),
                ({'name': 'Local Taste', 'location': 'Lviv', 'manager_id': 1}, 21, 404,
                 {'error': 'No such store_id 21'}),
                ({'name': 'Local Taste', 'location': 'Lviv', 'manager_id': 15}, 1, 404, {'error': 'No such user_id 15'})

        )
    )
    def test_update_store_by_id(self, update, adress, answer_code, answer):
        user = {'name': 'John Doe'}
        self.client.post('/users', json=user)
        newstore = {'name': 'Mad Cow', 'location': 'Lviv', 'manager_id': 1}
        self.client.post('/stores', json=newstore)
        resp = self.client.put(f'/stores/{adress}', json=update)
        assert resp.status_code == answer_code
        assert resp.json == answer


####################################################################
# Based on PartialConv2d
# BSD 3-Clause License
#
# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
#
# Author & Contact: Guilin Liu (guilinl@nvidia.com)
####################################################################

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

class ContentPartialConv2d(nn.Module):
    def __init__(self, *args, **kwargs):
        self.in_channels = kwargs['in_channels']
        self.out_channels = kwargs['out_channels']
        self.kernel_size = kwargs['kernel_size']
        self.stride = kwargs['stride']
        self.padding = kwargs['padding']
        self.dilation = kwargs.get('dilation', 1)
        self.groups = kwargs.get('groups', 1)
        self.multi_channel = kwargs.get('multi_channel', True)
        self.return_mask = kwargs.get('return_mask', True)
        self.bias = False
        self.mask_bias = None
        self.ratio = None

        super(ContentPartialConv2d, self).__init__()
        self.partialconv = nn.Conv2d(in_channels=self.in_channels,
                                     out_channels=self.out_channels,
                                     kernel_size=self.kernel_size,
                                     stride=self.stride,
                                     padding=self.padding,
                                     dilation=self.dilation,
                                     groups=self.groups)

    def update_mask(self, mask):
        if self.multi_channel:
            weight_maskUpdater = torch.ones([
                self.out_channels, self.in_channels, self.kernel_size,
                self.kernel_size
            ],
                                            device=mask.device)
        else:
            weight_maskUpdater = torch.ones(
                [1, 1, self.kernel_size, self.kernel_size], device=mask.device)

        update_mask = F.conv2d(mask,
                               weight_maskUpdater,
                               bias=self.mask_bias,
                               stride=self.stride,
                               padding=self.padding,
                               dilation=self.dilation,
                               groups=self.groups)
        update_mask = torch.clamp(update_mask, 0, 1)

        return update_mask

    def pretreat_input(self, input, mask):
        output1 = F.adaptive_avg_pool2d(input, 16)
        output2 = F.adaptive_avg_pool2d(input, 8)
        mask1 = F.adaptive_avg_pool2d(mask, 16)
        mask2 = F.adaptive_avg_pool2d(mask, 8)
        output1 = output1 / (mask1 + 1e-7)
        output2 = output2 / (mask2 + 1e-7)

        output1 = F.interpolate(output1,
                                size=[input.size(2)] * 2,
                                mode='nearest')
        output2 = F.interpolate(output2,
                                size=[input.size(2)] * 2,
                                mode='nearest')
        mask1 = F.interpolate(mask1, size=[input.size(2)] * 2, mode='nearest')
        mask2 = F.interpolate(mask2, size=[input.size(2)] * 2, mode='nearest')
        output1 = input + output1 * mask
        output2 = input + output2 * mask

        output1 = F.adaptive_avg_pool2d(output1, 16)
        output2 = F.adaptive_avg_pool2d(output2, 8)
        mask1 = F.adaptive_avg_pool2d(mask1, 16)
        mask2 = F.adaptive_avg_pool2d(mask1, 8)
        output1 = output1 / (mask1 + 1e-7)
        output2 = output2 / (mask2 + 1e-7)

        output1 = F.interpolate(output1,
                                size=[input.size(2)] * 2,
                                mode='bilinear',
                                align_corners=True)
        output2 = F.interpolate(output2,
                                size=[input.size(2)] * 2,
                                mode='bilinear',
                                align_corners=True)
        # self.ratio = nn.Parameter(torch.ones(1)).sigmoid()
        result = input + (0.7 * output1 + (0.3 * output2) * mask)
        return result

    def forward(self, input, mask):
        output = self.pretreat_input(input, mask)
        mask = self.update_mask(mask)
        output = self.partialconv(output)
        output = output * mask
        if self.return_mask:
            return output, mask
        else:
            return output


import math
import numpy as np
import matplotlib.pyplot as plt

def show_image(image):
    plt.cla()
    # plt.axis("off")
    if image.ndim == 2:
        plt.imshow(image, cmap="gray", vmin=0.0, vmax=1.0)
    else:
        plt.imshow(image)
    plt.show()

def show_boxes(image, boxes, boxes2=None, scores=None, show_axis=False):
    plt.cla()
    if not show_axis:
        plt.axis("off")

    if image.ndim == 2:
        plt.imshow(image, cmap="gray", vmin=0.0, vmax=1.0)
    elif image.shape[2] == 1:
        plt.imshow(image.squeeze(-1), cmap="gray", vmin=0.0, vmax=1.0)
    else:
        plt.imshow(image)

    for i, bbox in enumerate(boxes):
        # (x_min, y_min, w, h)
        rect = plt.Rectangle((bbox[0] - 0.5, bbox[1] - 0.5), bbox[2], bbox[3],
                             fill=False, edgecolor='red', linewidth=1.0)
        plt.gca().add_patch(rect)
        if scores is not None:
            plt.gca().text(bbox[0] - 0.5, bbox[1] - 0.5, str(scores[i]), color='red')

    # additional boxes, maybe GT
    if boxes2 is not None:
        for bbox2 in boxes2:
            rect2 = plt.Rectangle((bbox2[0] - 0.5, bbox2[1] - 0.5), bbox2[2], bbox2[3],
                                  fill=False, edgecolor='green', linewidth=1.0)
            plt.gca().add_patch(rect2)

    plt.show()

def draw_boxes(image, boxes, scores=None):
    """Draw bounding boxes on a gray image. Usually for tensorboard.

    Args:
        image (np.ndarray): (H, W, C)
        boxes (np.ndarray): (N, 4), represented as [x_min, y_min, w, h]
        scores (np.ndarray or list, optional): (N,)

    Returns:
        fig (plt.figure.Figure)

    """
    plt.close('all')
    fig = plt.figure()
    plt.axis("off")
    if image.ndim == 2:
        plt.imshow(image, vmin=0.0, vmax=1.0)
    elif image.shape[2] == 1:
        plt.imshow(image.squeeze(-1), cmap="gray", vmin=0.0, vmax=1.0)
    else:
        plt.imshow(image)
    for i, bbox in enumerate(boxes):
        rect = plt.Rectangle((bbox[0] - 0.5, bbox[1] - 0.5), bbox[2], bbox[3],
                             fill=False, edgecolor='red', linewidth=1.0)
        plt.gca().add_patch(rect)
        if scores is not None:
            plt.gca().text(bbox[0] - 0.5, bbox[1] - 0.5, '{:.2f}'.format(scores[i]), color='red')
    plt.tight_layout(0.02)
    return fig

def draw_images(images):
    """Draw gray images in a close-to-square layout

    Args:
        images (np.ndarray): (N, H, W) or (N, H, W, C)

    Returns:
        fig (plt.figure.Figure)

    """
    N, H, W = images.shape[:3]
    num_rows = int(math.sqrt(N))
    num_cols = (N + num_rows - 1) // num_rows

    plt.close('all')
    fig = plt.figure()
    for i in range(N):
        ax = fig.add_subplot(num_rows, num_cols, i + 1)
        image = images[i]
        if image.ndim == 2:
            ax.imshow(image, vmin=0.0, vmax=1.0)
        elif image.shape[2] == 1:
            ax.imshow(image.squeeze(-1), cmap="gray", vmin=0.0, vmax=1.0)
        else:
            ax.imshow(image)
        ax.set_axis_off()
    plt.subplots_adjust(left=0.02, right=0.98, bottom=0.05, top=0.95, wspace=0.01, hspace=0.01)
    return fig


#!/usr/bin/env python

"""

Web Application for the Catalog Project

Ensure you have run build_db.py and populate_db.py before running

"""

from flask import Flask, flash, request, redirect, url_for, render_template
from flask import jsonify, make_response
from flask import session as login_session

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

import random
import string
import requests

from oauth2client.client import flow_from_clientsecrets, FlowExchangeError

import httplib2

import json

from build_db import Items, Categories, Users, Base

CLIENT_ID = json.loads(open('secrets.json', 'r').read())['web']['client_id']

app = Flask(__name__)

engine = create_engine('sqlite:///catalog.db?check_same_thread=False')
Base.metadata.bind = engine

DBSession = sessionmaker(bind=engine)
session = DBSession()

""" Main Page / Displays Categories only"""

@app.route('/')
@app.route('/categories/')
def main():
    cat = session.query(Categories).all()
    print("Welcome to the main page!")

    # determine user privileges
    if 'email' not in login_session:
        user_id = "no value"
        client_logged_in = False
    else:
        user_id = getUserID(login_session['email'])
        client_logged_in = True
    is_admin = getUserAdminAccess(user_id)

    # render html based on user access
    return render_template(
        'main.html', categories=cat, is_admin=is_admin,
        client_logged_in=client_logged_in
    )

""" Items Page / Displays items once the user clicks on a category """

@app.route('/categories/<int:categories_id>/', methods=['GET', 'POST'])
def show_items(categories_id):
    print("This will show items by category")
    cat = session.query(Categories).filter_by(id=categories_id).one()
    item = session.query(Items).filter_by(category_id=categories_id).all()

    # determine user privileges
    if 'email' not in login_session:
        user_id = "no value"
        client_logged_in = False
    else:
        user_id = getUserID(login_session['email'])
        client_logged_in = True
    is_admin = getUserAdminAccess(user_id)

    # render html based on user access
    return render_template(
        'items.html', categories=cat, items=item,
        user_id=user_id, is_admin=is_admin,
        client_logged_in=client_logged_in
    )

""" New Items Page / allows user to add new item, or go back """

@app.route('/categories/<int:categories_id>/new/', methods=['GET', 'POST'])
def new_items(categories_id):
    print("This will allow a new item to be made")
    cat = session.query(Categories).filter_by(id=categories_id).one()
    return render_template('new_cat_item.html', category=cat)

@app.route('/categories/<int:categories_id>/new/entry/',
           methods=['GET', 'POST'])
def new_itm(categories_id):
    if request.method == 'POST':

        # Create item
        cat = session.query(Categories).filter_by(id=categories_id).one()
        user = session.query(Users).filter_by(id=getUserID(
            login_session['email']
        )).one()
        item_name = request.form['item_name']
        item_desc = request.form['item_desc']
        print("adding " + item_name + " to " + cat.name)
        new_item = Items(
            name=item_name, description=item_desc,
            owners=user, category=cat
        )

        # Add item to DB
        session.add(new_item)
        session.commit()
        cat = session.query(Categories).all()
        return render_template('main.html', categories=cat)
    else:
        Print("Something else happened... returning to main page")
        return render_template('main.html', categories=cat)

""" Remove items page / allows user to remove item, if they are the owner """

@app.route('/categories/<int:categories_id>/remove/<int:items_id>/',
           methods=['GET', 'POST'])
def remove_item(categories_id, items_id):
    print("This will allow the owner to remove an item")
    cat = session.query(Categories).filter_by(id=categories_id).one()
    itm = session.query(Items).filter_by(id=items_id).one()
    return render_template('remove_items.html', categories=cat, item=itm)

@app.route('/categories/<int:categories_id>/remove/<int:items_id>/confirm/',
           methods=['GET', 'POST'])
def remove_itm(categories_id, items_id):

    # Identify item for removal
    cat = session.query(Categories).filter_by(id=categories_id).one()
    itm = session.query(Items).filter_by(id=items_id).one()
    print("removing " + itm.name + " from " + cat.name)

    # remove item from DB
    session.delete(itm)
    session.commit()

    cat = session.query(Categories).all()
    return render_template('main.html', categories=cat)

""" Add new categories / Admin only """

@app.route('/categories/new/',  methods=['GET', 'POST'])
def new_category():
    print("Admin will be able to add category from here")
    return render_template('new_cat.html')

@app.route('/categories/new/entry/', methods=['GET', 'POST'])
def new_cat():
    if request.method == 'POST':

        # Create new category
        user_id = getUserID(login_session['email'])
        cat_name = request.form['category_name']
        print("adding: " + cat_name + " to categories")
        user = session.query(Users).filter_by(id=user_id).one()
        print("owner: " + user.name)
        category = Categories(name=cat_name, owner=user)

        # Add category to DB
        session.add(category)
        session.commit()
        cat = session.query(Categories).all()
        return render_template('main.html', categories=cat)
    else:
        return render_template('new_cat.html')

""" Remove categoryies / Admin only """

@app.route('/categories/remove/<int:categories_id>/',  methods=['GET', 'POST'])
def remove_category(categories_id):
    print("Admin will be able to remove category from here")
    cat = session.query(Categories).filter_by(id=categories_id).one()
    return render_template('remove.html', categories=cat)

@app.route('/categories/remove/<int:categories_id>/confirmed/',
           methods=['GET', 'POST'])
def remove(categories_id):

    # Identify Category
    cat = session.query(Categories).filter_by(id=categories_id).one()
    print("removing - " + cat.name)

    # Remove Category from DB
    session.delete(cat)
    session.commit()
    cat_list = session.query(Categories).all()
    return render_template('main.html', categories=cat_list)

# State token creation - taken from in-class examples

@app.route('/login/')
def show_login():
    state = ''.join(random.choice(string.ascii_uppercase + string.digits)
                    for x in xrange(32))
    login_session['state'] = state
    # return "session: %s" % login_session['state']
    return render_template('login.html', STATE=state)

# Google Connect

@app.route('/gconnect', methods=['POST'])
def gconnect():
    if request.args.get('state') != login_session['state']:
        response = make_response(json.dumps('Invalid state'), 401)
        response.headers['Content-Type'] = 'application/json'
        return response
    code = request.data

    try:
        oauth_flow = flow_from_clientsecrets('secrets.json', scope='')
        oauth_flow.redirect_uri = 'postmessage'
        credentials = oauth_flow.step2_exchange(code)
    except FlowExchangeError:
        response = make_response(json.dumps(
            'Failed to upgrade auth code'), 401)
        response.headers['Content-Type'] = 'application/json'
        return response

    # access token check
    access_token = credentials.access_token
    url = ('https://www.googleapis.com/oauth2/v1/tokeninfo?access_token=%s' %
           access_token)
    h = httplib2.Http()
    result = json.loads(h.request(url, 'GET')[1])
    if result.get('error') is not None:
        response = make_response(json.dumps(result.get('error')), 500)
        response.headers['Content-Type'] = 'application/json'
        return response

    # verify token is for intended user
    gplus_id = credentials.id_token['sub']
    if result['user_id'] != gplus_id:
        response = make_response(json.dumps('Token id mismatch'), 401)
        response.headers['Content-Type'] = 'application/json'
        return response

    # Verify token client id
    if result['issued_to'] != CLIENT_ID:
        response = make_response(json.dumps('Token Client ID mismatch'), 401)
        response.headers['Content-Type'] = 'application/json'
        return response

    # Check if user already logged in
    stored_access_token = login_session.get('credentials')
    stored_gplus_id = login_session.get('gplus_id')
    if stored_access_token is not None and gplus_id == stored_gplus_id:
        reponse = make_response(json.dumps('Already logged on'), 200)
        response.headers['Content-Type'] = 'application/json'
        return response

    # Store access token in session for later use
    login_session['access_token'] = credentials.access_token
    login_session['gplus_id'] = gplus_id

    # Get user info
    user_info_url = "https://www.googleapis.com/oauth2/v1/userinfo"
    params = {'access_token': credentials.access_token, 'alt': 'json'}
    answer = requests.get(user_info_url, params=params)

    data = answer.json()

    login_session['username'] = data['name']
    login_session['picture'] = data['picture']
    login_session['email'] = data['email']

    # add provider
    login_session['provider'] = 'google'

    # check if user exists / prompt account creation
    user_id = getUserID(data["email"])
    if not user_id:
        user_id = createUser(login_session)
    login_session['user_id'] = user_id

    output = ''
    output += '
 Hi, ' + login_session['username'] + '
'
    output += ''
    flash("You are logged in as: %s" % login_session['username'])

    return output

@app.route('/logoff/')
def sign_out():

    access_token = login_session.get('access_token')

    url = 'https://accounts.google.com/o/oauth2/revoke?token=%s' % access_token
    h = httplib2.Http()
    result = h.request(url, 'GET')[0]

    login_session.clear()

    output = ''
    output += '
 Signing Out, Please Wait
'
    output += ''
    return output

""" Helper functions: Create User, Get User Info/User ID/User Access Rights"""

def createUser(login_session):
    newUser = Users(name=login_session['username'], email=login_session[
        'email'], picture=login_session['picture'], admin=False)
    session.add(newUser)
    session.commit()
    user = session.query(Users).filter_by(email=login_session['email']).one()
    return user.id

def getUserInfo(user_id):
    user = session.query(Users).filter_by(id=user_id).one()
    return user

def getUserID(email):
    try:
        user = session.query(Users).filter_by(email=email).one()
        return user.id
    except:
        return None

def getUserAdminAccess(user_id):
    try:
        user = session.query(Users).filter_by(id=user_id).one()
        if user.admin:
            return True
        else:
            return False
    except:
        print("No user found")
        return False

""" JSON Endpoints """

# List all items in a category
@app.route('/categories/<int:categories_id>/JSON/')
def itemInCategoryJSON(categories_id):
    items = session.query(Items).filter_by(category_id=categories_id).all()
    return jsonify(items=[i.serialize for i in items])

# List a specific item
@app.route('/categories/<int:categories_id>/items/<int:items_id>/JSON/')
def itemJSON(categories_id, items_id):
    categories = session.query(Categories).filter_by(id=categories_id).one()
    items = session.query(Items).filter_by(
        id=items_id, category_id=categories.id).one()
    return jsonify(items=[items.serialize])

# List all categories
@app.route('/categories/JSON/')
def categoryJSON():
    categories = session.query(Categories).all()
    return jsonify(categories=[c.serialize for c in categories])

if __name__ == '__main__':
    app.secret_key = 'in8_Zp8k5yLKpA-LUFJ0unNi'
    app.debug = True
    app.run(host='0.0.0.0', port=5000)


from numpy import array, dstack, roll
from skimage.transform import rescale
from numpy import int64
from numpy import square

def MSE(image1, image2): #очень хорошо совмещает
	im1 = array(image1, dtype=int64)
	im2 = array(image2, dtype=int64)
	return (square(im1 - im2)).sum().mean(axis=None)

def allocate_channels(bgr_image):
	height = len(bgr_image)
	b_image = bgr_image[:height / 3]
	g_image = bgr_image[height / 3: 2 * (height / 3)]
	r_image = bgr_image[2 * (height / 3):]
	height = min(len(b_image), len(g_image), len(r_image))
	b_image = b_image[:height]
	r_image = r_image[:height]
	g_image = g_image[:height]
	width = len(b_image[0])
	magic_number = 0.05
	b_image = b_image[height * magic_number: -height * magic_number, width * magic_number: -width * magic_number]
	g_image = g_image[height * magic_number: -height * magic_number, width * magic_number: -width * magic_number]
	r_image = r_image[height * magic_number: -height * magic_number, width * magic_number: -width * magic_number]
	return (b_image, g_image, r_image)

def find_best_shift(image1, image2):
	height = len(image1)
	width = len(image1[0])
	critical_size = 200
	if width > critical_size:
		basic_shift = find_best_shift(rescale(image1, 0.5, preserve_range=True),
			rescale(image2, 0.5, preserve_range = True))
		basic_shift = [c * 2 for c in basic_shift]
		max_shift = 1
	else:
		max_shift = 10
		basic_shift = [0, 0]
	best_shift = [0, 0]
	best_result = MSE(image1, image2)
	for x_shift_iter in range(-max_shift, max_shift + 1):
		for y_shift_iter in range(-max_shift, max_shift + 1):
			x_shift = x_shift_iter + basic_shift[0]
			y_shift = y_shift_iter + basic_shift[1]
			left = max(0, x_shift)
			right = min(width, width + x_shift)
			top = max(0, y_shift)
			bottom = min(height, height + y_shift)
			current_result = MSE(image1[top:bottom, left:right],
				image2[top - y_shift:bottom - y_shift,left - x_shift:right - x_shift])
			if best_result > current_result:
				best_shift = (x_shift, y_shift)
				best_result = current_result
	return best_shift

def overlay_channels(blue, blue_coords, green, green_coords, red, red_coords):
	height = len(blue)
	width = len(blue[0])
	left = max(blue_coords[0], green_coords[0], red_coords[0])
	right = min(width + blue_coords[0], width + green_coords[0], width + red_coords[0])
	top = max(blue_coords[1], green_coords[1], red_coords[1])
	bottom = min(height + blue_coords[1], height + green_coords[1], height + red_coords[1])
	blue = blue[top - blue_coords[1]:bottom - blue_coords[1],
	left - blue_coords[0]:right - blue_coords[0]]
	green = green[top - green_coords[1]:bottom - green_coords[1],
	left - green_coords[0]:right - green_coords[0]]
	red = red[top - red_coords[1]:bottom - red_coords[1],
	left - red_coords[0]:right - red_coords[0]]
	return dstack((red, green, blue))

def align(bgr_image):
	b_image, g_image, r_image = allocate_channels(bgr_image)
	gb_x, gb_y = find_best_shift(g_image, b_image)
	gr_x, gr_y = find_best_shift(g_image, r_image)

	return overlay_channels(b_image, (gb_x, gb_y), g_image, (0, 0), r_image, (gr_x, gr_y))


from .types import DataType
from tensorflow.core.framework.tensor_pb2 import TensorProto
from tensorflow.core.framework.tensor_shape_pb2 import TensorShapeProto
import numpy as np
from typing import Iterable, AnyStr, Tuple

def coerce_to_bytes(text: AnyStr) -> bytes:
    if isinstance(text, str):
        return text.encode("utf-8")
    else:
        return text

def write_values_to_tensor_proto(
    tensor_proto: TensorProto, values: Iterable, dtype: DataType
) -> TensorProto:
    proto_field = getattr(tensor_proto, f"{dtype.proto_field_name}")
    if dtype.is_numeric:
        proto_field.extend(values)
    else:
        proto_field.extend([coerce_to_bytes(v) for v in values])
    tensor_proto.tensor_content = values.tobytes()
    return tensor_proto

def ndarray_to_tensor_proto(ndarray: np.ndarray) -> TensorProto:
    dtype = DataType(ndarray.dtype.type)
    proto = TensorProto(
        dtype=dtype.enum,
        tensor_shape=TensorShapeProto(dim=[TensorShapeProto.Dim(size=d) for d in ndarray.shape]),
    )
    proto = write_values_to_tensor_proto(tensor_proto=proto, values=ndarray.ravel(), dtype=dtype)
    return proto

def extract_shape(tensor_proto: TensorProto) -> Tuple[int, ...]:
    return tuple((int(d.size) for d in tensor_proto.tensor_shape.dim))

def tensor_proto_to_ndarray(tensor_proto: TensorProto) -> np.ndarray:
    dtype = DataType(tensor_proto.dtype)
    shape = extract_shape(tensor_proto)
    if tensor_proto.tensor_content:
        return (np.frombuffer(tensor_proto.tensor_content,
                              dtype=dtype.numpy_dtype).copy().reshape(shape))
    proto_values = getattr(tensor_proto, dtype.proto_field_name)
    return np.array(proto_values, dtype=dtype.numpy_dtype).reshape(*shape)


# coding=utf-8

from flask import Flask, render_template, request, make_response
from flask import jsonify

import time
import threading
import demo as demo

"""
定义心跳检测函数

"""

def heartbeat():
    print(time.strftime('%Y-%m-%d %H:%M:%S - heartbeat', time.localtime(time.time())))
    timer = threading.Timer(60, heartbeat)
    timer.start()

timer = threading.Timer(60, heartbeat)
timer.start()

"""
ElementTree在 Python 标准库中有两种实现。
一种是纯 Python 实现例如 xml.etree.ElementTree ，
另外一种是速度快一点的 xml.etree.cElementTree 。
 尽量使用 C 语言实现的那种，因为它速度更快，而且消耗的内存更少
"""

app = Flask(__name__, static_url_path="/static")

@app.route('/message', methods=['POST'])
# """定义应答函数，用于获取输入信息并返回相应的答案"""
def reply():
    # 从请求中获取参数信息
    req_msg = request.form['msg']
    if req_msg == 'train()':
        # 执行训练
        res_msg = demo.train()
        res_msg = res_msg.replace(' ', '')
        return jsonify({'text': res_msg})
    res_msg = demo.predict(req_msg)
    # 将语句使用结巴分词进行分词
    # req_msg =" ".join(jieba.cut(req_msg))
    # 调用decode_line对生成回答信息
    # res_msg = execute.decoder_online(sess, conf.gen_config ,model, vocab ,rev_vocab, req_msg )
    # res_msg=generator.get_predicted_sentence(sess,req_msg,vocab,model,conf.gen_config.beam_size,conf.gen_config.buckets)
    # 将unk值的词用微笑符号袋贴
    # res_msg = res_msg.replace('_UNK', '^_^')
    res_msg = res_msg.replace(' ','')
    return jsonify({'text': res_msg})

"""
jsonify:是用于处理序列化json数据的函数，就是将数据组装成json格式返回

http://flask.pocoo.org/docs/0.12/api/#module-flask.json
"""

@app.route("/")
def index():
    return render_template("index.html")

#
'''
初始化seq2seqModel，并进行动作

    1. 调用执行器的主程序
    2. 生成一个在线decode进程，来提供在线聊天服务
'''
# _________________________________________________________________

# _________________________________________________________________

# 启动APP
if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8808)


import numpy as np
import pandas as pd

seed = 123
np.random.seed(seed)

# X: Independent Variables (player/team stats), y: Dependent (actual runs scored)
X = pd.read_csv("MLB/X_sort.csv", encoding='latin-1')
y = pd.read_csv("MLB/y_sort.csv", encoding='latin-1', names=['Score'])

from constant_variables import features_top_list
X2 = X['Expected_Runs']
X = X.drop('Expected_Runs', axis = 1)

#Splitting the data into a training set and a test set(determine accuracy of the algo on data it hasn't seen)
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size = 0.20, random_state = 0)
'''
# Feature Scaling (Important for high intensity computations)
from sklearn.preprocessing import StandardScaler
sc = StandardScaler()
X_train = sc.fit_transform(X_train)
X_test = sc.transform(X_test)

# Applying PCA
# explained_variance - list of principle components and % of variance explained by each of them
# Run with n_components = None 1st (look at explained_variance)
from sklearn.decomposition import PCA
pca = PCA(n_components = 60)
X_train = pca.fit_transform(X_train)
X_test = pca.transform(X_test)
explained_variance = pca.explained_variance_ratio_
'''

#Training the regression algorithm on the training set (It is basically trying to determine the importance of each variable )
# Regression sums up the effect of the variables and spits out a prediction.  It's a continuous variable.
# Fitting XGBoost to the Training set
from xgboost import XGBRegressor
classifier = XGBRegressor()
classifier.fit(X_train, y_train)

#Training the classification algorithm on the training set (It is basically trying to determine the importance of each variable )
#In classification, you are trying to group things into neat catagories.  ex. Is it a monkey or a dog (discrete variable)
# Fitting XGBoost to the Training set
from xgboost import XGBClassifier
classifier1 = XGBClassifier()
classifier1.fit(X_train, y_train)

# 3 methods of evaluating the prediction quality (mean absolute error, root mean squared error, correlation coefficient)
def eval_metrics(y_test, y_pred):
    mae = np.round(abs(y_test.values - y_pred).mean(),decimals=2)
    rmse = (((y_test.values - y_pred)**2).mean())**(1/2)
    r = (((y_test.values-y_test.values.mean())*(y_pred - y_pred.mean())).sum()) / ((1-len(y_test))*(y_test.values.std())*(y_pred.std()))
    print('mae: {0:.2f} \nrmse: {1:.2f} \nr: {2:.4f}'.format(mae, rmse, r))
    return mae, rmse, r

# Predicting the Test set results
y_pred_X = classifier.predict(X_train)
y_pred_X_test = classifier.predict(X_test)
y_pred_X_test_rd = np.round(y_pred_X_test)

# Predicting the Test set results
y_pred_cls = classifier1.predict(X_train)
y_pred_cls_test = classifier1.predict(X_test)
y_pred_cls_test_rd = np.round(y_pred_cls_test)

metrics = eval_metrics(y_test, y_pred_X_test_rd)
metrics_class = eval_metrics(y_test, y_pred_cls_test_rd)

# A method of testing your algorithm that splits data into small segments(batches) and averaging the mean of their accuracies
# Applying k-Fold Cross Validation
from sklearn.model_selection import cross_val_score
accuracies = cross_val_score(estimator = classifier, X = X_train, y = y_train, cv = 10)
mean_acc = accuracies.mean()
train_std = accuracies.std()
classifier.evals_result

# What variables are most important?  Higher the #, the greater the importance
feature_importance = classifier1.feature_importances_
#features = pd.Series(feature_importance, index = len(X_train[0]))
#features_sorted = pd.Series(features.sort_values())

# This is a neural network
# It runs through the data 2000 times adjusting the weights every 35 games
# The dropout is used to prevent overfitting - training too closely to the training data at the expense of test accuracy
# It drops out variables randomly
# Output dim is the number of nodes in each layer
# The goal is to reduce a cost function as much as possible (mean squared error)
import keras
from keras import optimizers
from keras.models import Sequential
from keras.layers import Dense, Activation
from keras.layers import Dropout
from keras import objectives
from keras import backend as K

# Initialising the ANN
regressor1 = Sequential()

# Adding the input layer and the first hidden layer with dropout (start 0.1 (10%))
regressor1.add(Dense(output_dim = 90, init = 'uniform', activation = 'relu', input_dim = len(X_train.columns)))
regressor1.add(Dropout(p=0.1))

# Adding the second hidden layer
regressor1.add(Dense(output_dim = 75, init = 'uniform', activation = 'relu'))
regressor1.add(Dropout(p=0.1))

# Adding the third hidden layer
regressor1.add(Dense(output_dim = 60, init = 'uniform', activation = 'relu'))
regressor1.add(Dropout(p=0.1))

# Adding the 4th hidden layer
regressor1.add(Dense(output_dim = 45, init = 'uniform', activation = 'relu'))
regressor1.add(Dropout(p=0.1))

# Adding the output layer
regressor1.add(Dense(output_dim = 1, init = 'uniform'))

# Compiling the ANN (adam[SGD] - optimizer function to find optimal weights)
# Binary Dept Var(Binary_CrossEntropy) Dependent Var > 2 Outcomes (Categorical_CrossEntropy)
# [Accuracy] in brackets because list expected
regressor1.compile(optimizer = 'adam', loss = 'mse', metrics=['binary_crossentropy','acc'])

# Fitting the ANN to the Training set
regressor1.fit(X_train, y_train, batch_size = 35, epochs = 2000)

# Predicting the Test set results
y_pred_test = regressor1.predict(X_test)
y_pred_test_rd = np.round(y_pred_test)
y_test_index = [x for x in y_test.index]

# Predicting a new result
y_pred_train = regressor1.predict(X_train)
y_pred_train_rd = np.round(y_pred_train)
y_train_index = [x for x in y.index]

y_test_total_index = [x for x in pd.Series(y_test.index)]

expected = X2[:]
y_expected_test = expected[y_test_total_index]
y_expected_test_round = np.round(y_expected_test)

metrics_reg = eval_metrics(y_test, y_pred_test)
metrics_expect = eval_metrics(y_test, y_expected_test_round.values)

y_test_total = ((pd.DataFrame(y_pred_test_rd) + pd.DataFrame(y_pred_X_test_rd.ravel()))/2).astype(int)
metrics_total = eval_metrics(y_test, y_test_total[0].values)

from sklearn.neural_network import MLPRegressor
regressor2 = MLPRegressor(solver='adam', max_iter=1000, random_state=9)
regressor2.fit(X_train, y_train)
y_pred_mlp = regressor2.predict(X_test)
y_pred_mlp_rd = np.round(y_pred_mlp)

for i,x in enumerate(y_pred_mlp_rd):
    if x<0:
        y_pred_mlp_rd[i] = 0

#y_pred_mlp_rd = [0 for i,x in enumerate(y_pred_mlp_rd) if x<0]

test_eval = eval_metrics(y_test, y_pred_mlp_rd)

y_test_total2 = ((pd.DataFrame(y_pred_test_rd) + pd.DataFrame(y_pred_X_test_rd.ravel()) + pd.DataFrame(y_pred_mlp_rd.ravel()))/3).astype(int)
metrics_total2 = eval_metrics(y_test, y_test_total2[0].values)

from catboost import Pool, CatBoostRegressor, cv
from sklearn.metrics import accuracy_score
model = CatBoostRegressor()

# Fit model
model.fit(
    X_train, y_train,
    eval_set=(X_test, y_test),
#     logging_level='Verbose',  # you can uncomment this for text output
    plot=True
)

model.fit(X_train, y_train, plot=True)
# Get predictions
pred_cat = model.predict(X_test)
pred_cat_rd = np.round(pred_cat)
cat_eval = eval_metrics(y_test, pred_cat_rd)

y_test_total3 = ((pd.DataFrame(y_pred_test_rd) + pd.DataFrame(y_pred_X_test_rd.ravel()) + pd.DataFrame(y_pred_mlp_rd.ravel()) + pd.DataFrame(pred_cat_rd.ravel()))/4).astype(int)
metrics_total3 = eval_metrics(y_test, y_test_total3[0].values)

model.feature_importances_

params = {'iterations': 500, 'depth': 4, 'loss_function': 'MAE', 'eval_metric':'RMSE', 'learning_rate':0.1, 'logging_level': 'Silent'}
pool = Pool(X_train, y_train)
#scores = cv(params, pool)

metrics_expect
est1, est2, est3, est1s, est2s, est3s = [], [], [], [], [], []
estimators = [metrics, metrics_class, metrics_reg, metrics_total, test_eval, metrics_total2, cat_eval, metrics_total3]
estimator_names = ['metrics', 'metrics_class', 'metrics_reg', 'metrics_total', 'test_eval', 'metrics_total2', 'cat_eval', 'metrics_total3']
for i,est in enumerate(estimators):
    est1.append(est[0])
    est1s.append(estimator_names[i])
    est2.append(est[1])
    est2s.append(estimator_names[i])
    est3.append(est[2])
    est3s.append(estimator_names[i])
mae_total = pd.Series(est1, index = est1s, name = 'mae').sort_values()
rsme_total = pd.Series(est2, index = est2s, name = 'rsme').sort_values()
r_total = pd.Series(est3, index = est3s, name = 'r').sort_values()

y_cat_met = ((pd.DataFrame(y_pred_X_test_rd.ravel()) + pd.DataFrame(pred_cat_rd.ravel()))/2).astype(int)
cat_metrics = eval_metrics(y_test, y_cat_met[0].values)

y_cat_met_cls = ((pd.DataFrame(y_pred_X_test_rd.ravel()) + pd.DataFrame(pred_cat_rd.ravel()) + pd.DataFrame(y_pred_cls_test_rd.ravel()))/3).astype(int)
cat_metrics_cls = eval_metrics(y_test, y_cat_met_cls[0].values)

y_cat_met_cls_ratio = (pd.DataFrame((y_pred_X_test_rd*.4).ravel()) + pd.DataFrame((pred_cat_rd*.5).ravel()) + pd.DataFrame((y_pred_cls_test_rd.astype(float)*.1).ravel())).astype(int)
cat_metrics_cls_ratio = eval_metrics(y_test, y_cat_met_cls_ratio[0].values)

y_cat_met_cls_ratio = (pd.DataFrame((y_pred_X_test*.4).ravel()) + pd.DataFrame((pred_cat*.5).ravel()) + pd.DataFrame((y_pred_cls_test.astype(float)*.1).ravel())).astype(int)
cat_metrics_cls_ratio = eval_metrics(y_test, y_cat_met_cls_ratio[0].values)


from Actor import Actor
from Critic import Critic
from Noise import Noise
from ReplayBuffer import ReplayBuffer

import numpy as np
import random

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

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

class Agent():
    """DDQN Agent utilzing an actor/critic nn. """

    def __init__(self, state_size, action_size, memory, id, seed=0):
        """Initialize Agent
        Args:
            state_size: number of possible states
            action_size: number of possible acitons
            seed: random seed
        """
        self.state_size = state_size
        self.action_size = action_size
        self.epsilon = 1
        self.discount_factor = 1e-3
        self.t_step = 0
        self.seed = random.seed(seed)
        self.id = id

        # Actor model
        self.actor_local = Actor(state_size, action_size, seed).to(device)
        self.actor_target = Actor(state_size, action_size, seed).to(device)
        self.actor_optimizer = optim.Adam(self.actor_local.parameters(), lr=1e-3) # Learning rate

        # Critic Model
        self.critic_local = Critic(state_size, action_size, seed).to(device)
        self.critic_target = Critic(state_size, action_size, seed).to(device)
        self.critic_optimizer = optim.Adam(self.critic_local.parameters(), lr=1e-3, weight_decay=0) # Learning rate

        # Add noise
        self.noise = Noise(action_size, seed)

        # Memory buffer
        self.memory = memory

    def reset_noise(self):
        """ Reset noise"""
        self.noise.reset()

    def step(self, state, action, reward, next_state, done, step):
        """Save the experience in the buffer
        Args:
            state: current state
            action: selected action
            reward: reward given
            next_state: next state to advance to
            done: has the memory completed
        """
        # save experience
        self.memory.add(state, action, reward, next_state, done)
        # Learn in batches
        if len(self.memory) > self.memory.batch_size:
            experiences = self.memory.get_random_sample()
            self.learn(experiences, 0.99) # discount factor

    def learn(self, experiences, gamma):
        """Teachs the agent based on the current state, action reward and next state

        Args:
            experiences: tuple of state, actions, rewards, next states and done
            gamma: the discount factor
        """
        states, actions, rewards, next_states, dones = experiences

        # Update critic network
        # Get next state-action and Q values
        actions_next = self.actor_target(next_states)
        if self.id == 0:
            actions_next = torch.cat((actions_next, actions[:,2:]), dim=1)
        else:
            actions_next = torch.cat((actions[:,:2], actions_next), dim=1)

        Q_targets_next = self.critic_target(next_states, actions_next)
        # Compute Q targets for current states
        Q_targets = rewards + (gamma * Q_targets_next * (1 - dones))
        # Get expected Q values from local model
        Q_expected = self.critic_local(states, actions)
        critic_loss = F.mse_loss(Q_expected, Q_targets)
        # Minimize loss
        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        self.critic_optimizer.step()

        # Update actor network
        actions_pred = self.actor_local(states)
        if self.id == 0:
            actions_pred = torch.cat((actions_pred, actions[:,2:]), dim=1)
        else:
            actions_pred = torch.cat((actions[:,:2], actions_pred), dim=1)

        actor_loss = -self.critic_local(states, actions_pred).mean()
        # minimize loss
        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        self.actor_optimizer.step()

        # Updates networks
        self.update(self.critic_local, self.critic_target, 6e-2)
        self.update(self.actor_local, self.actor_target, 6e-2)

        # reduce epsilon per learned event
        self.epsilon -= self.discount_factor
        self.epsilon = max(self.epsilon, 0)
        # Reset noise between learning events
        self.noise.reset()

    def update(self, local_model, target_model, tau):
        """Update the model params
        target = t*local+(1-t)*target

        Args:
            local_model: pytorch model source
            target_model: pytorch model destination
            tau: interpolation param
        """
        for target_param, local_param in zip(target_model.parameters(), local_model.parameters()):
            target_param.data.copy_(tau*local_param.data + (1.0-tau)*target_param.data)

    def get_action(self, state, add_noise=True):
        """Gets the action from the given state
        Args:
            state: current state
            eps: epsilon value for a epsilon-greedy policy
        Returns:
            action selection based on the epsilon-greedy policy
        """
        state = torch.from_numpy(state).float().to(device)
        self.actor_local.eval()

        with torch.no_grad():
            action = self.actor_local(state).cpu().data.numpy()
        self.actor_local.train()
        # Add noise discounted by epsilon
        if add_noise:
            action += self.epsilon * self.noise.sample()
        return np.clip(action, -1, 1)