import dash
import dash_core_components as dcc
import dash_html_components as html
from dash.dependencies import Input, Output, State
from dash.exceptions import PreventUpdate
import pandas as pd
import pybaseball as pb
import plotly.express as px

external_stylesheets = ['https://stackpath.bootstrapcdn.com/bootstrap/4.4.1/css/bootstrap.min.css']

app = dash.Dash(__name__, external_stylesheets=external_stylesheets)
server = app.server

app.layout = html.Div(children=[
  dcc.Markdown(children="## Plotly Express Data Sets"),
  html.Br(),
  html.Br(),
  dcc.Markdown(children="### Plotly Express comes with several data sets built-in."),
  dcc.Dropdown(
    id="dropdown",
    options=[
    {'label': 'gapminder', 'value': 'gapminder'},
    {'label': 'tips', 'value': 'tips'},
    {'label': 'iris', 'value': 'iris'},
    {'label': 'wind', 'value': 'wind'},
    {'label': 'election', 'value': 'election'},
    {'label': 'carshare', 'value': 'carshare'}],
    value='gapminder'),
  dcc.Markdown(id='blurb'),
  dcc.Graph(id='graph'),
])

@app.callback(
    [Output(component_id='blurb', component_property='children'),
            Output(component_id='graph', component_property='figure')],
    [Input('dropdown', 'value')])
def update_output_graph(value):
    if value == "gapminder":
        blurb = """ Each row represents a country on a given year.
            https://www.gapminder.org/data/
            Returns:
                A `pandas.DataFrame` with 1704 rows and the following columns: `['country', 'continent', 'year', 'lifeExp', 'pop', 'gdpPercap',
            'iso_alpha', 'iso_num']`.
        """
        df = px.data.gapminder()
        fig = px.scatter(df, x="gdpPercap", y="lifeExp",             animation_frame="year", animation_group="country",
           size="pop", color="continent", hover_name="country", facet_col="continent",
           log_x=True, size_max=45, range_x=[100,100000], range_y=[25,90])
        fig.update_layout(title="gdp per cap vs life expectancy over time")

    if value == "iris":
        blurb = """
            Each row represents a flower.
            https://en.wikipedia.org/wiki/Iris_flower_data_set
            Returns:
                A `pandas.DataFrame` with 150 rows and the following columns: `['sepal_length', 'sepal_width', 'petal_length', 'petal_width', 'species',
            'species_id']`.
        """
        df = px.data.iris()
        fig = px.density_contour(df, x="sepal_width", y="sepal_length", color="species", marginal_x="rug", marginal_y="histogram")

    if value == "tips":
        blurb =  """
            Each row represents a restaurant bill.
            https://vincentarelbundock.github.io/Rdatasets/doc/reshape2/tips.html
            Returns:
                A `pandas.DataFrame` with 244 rows and the following columns: `['total_bill', 'tip', 'sex', 'smoker', 'day', 'time', 'size']`.
        """
        df = px.data.tips()
        fig = px.bar(df, x="sex", y="total_bill", color="smoker", barmode="group", facet_row="time", facet_col="day",
            category_orders={"day": ["Thur", "Fri", "Sat", "Sun"], "time": ["Lunch", "Dinner"]})
        fig.update_layout(title="tips by gender, day of week, and mealtime")

    if value == "election":
        blurb = """
            Each row represents voting results for an electoral district in the 2013 Montreal mayoral election.
            Returns:
                A `pandas.DataFrame` with 58 rows and the following columns: `['district', 'Coderre', 'Bergeron', 'Joly', 'total', 'winner', 'result']`.
        """
        df = px.data.election()
        fig = px.scatter_ternary(df, a="Joly", b="Coderre", c="Bergeron", color="winner", size="total", hover_name="district", size_max=15, color_discrete_map = {"Joly": "blue", "Bergeron": "green", "Coderre":"red"} )

    if value == "wind":
        blurb =  """
            Each row represents a level of wind intensity in a cardinal direction, and its frequency.
            Returns:
                A `pandas.DataFrame` with 128 rows and the following columns: `['direction', 'strength', 'frequency']`.
        """
        df = px.data.wind()
        fig = px.scatter_polar(df, r="frequency", theta="direction", color="strength", symbol="strength",
                    color_discrete_sequence=px.colors.sequential.Plasma_r)

    if value == "carshare":
        blurb = """
            Each row represents the availability of car-sharing services near the centroid of a zone in Montreal.
            Returns:
                A `pandas.DataFrame` with 249 rows and the following columns: `['centroid_lat', 'centroid_lon', 'car_hours', 'peak_hour']`.
        """
        df = px.data.carshare()
        fig = px.scatter_mapbox(df, lat="centroid_lat", lon="centroid_lon", color="peak_hour", size="car_hours",
        color_continuous_scale=px.colors.cyclical.IceFire, size_max=15, zoom=10, mapbox_style="open-street-map")

    fig.update_yaxes(automargin=True)

    return blurb, fig

if __name__ == '__main__':
    app.run_server(debug=True)


#!/usr/bin/env python

import sys
import re
from time import sleep
from datetime import datetime

import traceback
import usb.core
import argparse
import struct
import sys

HOST_TO_DEVICE = 0x40
DEVICE_TO_HOST = 0xC0
TIMEOUT_MS = 1000

MAX_PAGES = 8
PAGE_SIZE = 64
EEPROM_FORMAT = 8

# An extensible, stateful "Test Fixture"
class Fixture(object):
    def __init__(self):
        self.eeprom_pages = None
        self.subformat = None
        self.dev = None

        parser = argparse.ArgumentParser(
            description="Command-line utility to play with lasers (INHERENTLY DANGEROUS!)",
            epilog="To manually turn the laser off after an arbitrary delay, use --max-ms 0. " +
                   "To exit the script with laser still firing (HIGHLY DANGEROUS), use --max-ms -1."
        )
        parser.add_argument("--pid",                    default="4000",      help="USB PID in hex (default 4000)", choices=["1000", "2000", "4000"])
        parser.add_argument("--debug",                  action="store_true", help="debug output")
        parser.add_argument("--enable",                 action="store_true", help="fire laser")
        parser.add_argument("--max-ms",                 type=int,            help="firing time (ms) (default 1000)", default=1000)
        parser.add_argument("--mod-enable",             action="store_true", help="enable modulation")
        parser.add_argument("--mod-period-us",          type=int,            help="laser modulation pulse period (us) (default 1000)", default=1000)
        parser.add_argument("--mod-width-us",           type=int,            help="laser modulation pulse width (us) (default 100)", default=100)

        self.args = parser.parse_args()

        # convert PID from hex string
        self.pid = int(self.args.pid, 16)

        # find the FIRST connected spectrometer of the given PID
        self.dev = usb.core.find(idVendor=0x24aa, idProduct=self.pid)
        if not self.dev:
            print("No spectrometers found with PID 0x%04x" % self.pid)

    def run(self):

        # handle disable operation first
        if not self.args.enable:
            self.set_enable(False)
            return

        # apparently we're to fire the laser

        if self.args.mod_enable:
            self.set_modulation_params()
        else:
            self.set_modulation_enable(False)

        self.set_enable(True)

        if self.args.max_ms > 0:
            self.sleep_ms(self.args.max_ms)
        elif self.args.max_ms == 0:
            cont = input("\nPress <enter> to disable laser...")
        elif self.args.max_ms == -1:
            print("DANGER -- Exiting with laser still firing!!!")
            sys.exit(1)

        self.set_enable(False)

    def sleep_ms(self, ms):
        print(f"sleeping {ms} ms...")
        if ms > 1000 and self.pid == 0x4000:
            # monitor battery while sleeping
            start = datetime.now()
            while (datetime.now() - start).total_seconds() * 1000.0 < ms:
                print("Battery: " + self.get_battery_level())
                sleep(1)
        else:
            sleep(ms/1000.0)

    ############################################################################
    # opcodes
    ############################################################################

    def get_battery_level(self):
        raw = self.get_cmd(0xff, 0x13, 3)
        percentage = raw[1] + (1.0 * raw[0] / 256.0)
        charging = raw[2] != 0
        return f"%s (%.2f%%) (%s)" % (raw, percentage, "charging" if charging else "not charging")

    def set_enable(self, flag):
        print("setting LASER_ENABLE %s" % ("on" if flag else "off"))
        self.send_cmd(0xbe, 1 if flag else 0)

    def set_modulation_enable(self, flag):
        print("setting LASER_MOD_ENABLE %s" % ("on" if flag else "off"))
        self.send_cmd(0xbd, 1 if flag else 0)

    def set_modulation_params(self):
        if self.args.mod_period_us > 0xffff or \
           self.args.mod_width_us > 0xffff:
            print("error: lame script doesn't support full 40-bit 5-byte args")
            return

        # should we modulate after all?
        if self.args.mod_period_us <= self.args.mod_width_us:
            print("disabling modulation because period %d <= width %d" % (self.args.mod_period_us, self.args.mod_width_us))
            self.set_modulation_enable(False)
            return

        print("setting LASER_MOD_PULSE_PERIOD %d" % self.args.mod_period_us)
        self.send_cmd(0xc7, self.args.mod_period_us, buf=[0]*8)

        print("setting LASER_MOD_PULSE_WIDTH %d" % self.args.mod_width_us)
        self.send_cmd(0xdb, self.args.mod_width_us, buf=[0]*8)

        self.set_modulation_enable(True)

    ############################################################################
    # Utility Methods
    ############################################################################

    def is_arm(self):
        return self.pid == 0x4000

    def is_sig(self):
        return self.is_arm() # close enough

    def str2bool(self, s):
        return s.lower() in ("true", "yes", "on", "enable", "1")

    def debug(self, msg):
        if self.args.debug:
            print("DEBUG: %s" % msg)

    def send_cmd(self, cmd, value, index=0, buf=None):
        if buf is None:
            if self.is_arm():
                buf = [0] * 8
            else:
                buf = 0
        self.debug("ctrl_transfer(0x%02x, 0x%02x, 0x%04x, 0x%04x) >> %s" % (HOST_TO_DEVICE, cmd, value, index, buf))
        self.dev.ctrl_transfer(HOST_TO_DEVICE, cmd, value, index, buf, TIMEOUT_MS)

    def get_cmd(self, cmd, value=0, index=0, length=64):
        return self.dev.ctrl_transfer(DEVICE_TO_HOST, cmd, value, index, length, TIMEOUT_MS)

fixture = Fixture()
if fixture.dev:
    fixture.run()


import socket

target_host = 'jpandrade.info'

print('[+] Creating socket...')
client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
print('[!] Socket created!')
print('[+] Connecting with remote host...')
client.connect((target_host, 80))
print('[+] Connection ok')
request = 'GET / HTTP/1.1\r\nHost:%s\r\n\r\n' % target_host
client.send(request.encode())
while True:
    data = client.recv(128)
    print(data.decode())
    if data == '':
        break
print('[-] Closing the socket')
client.close()


from django.db.models.signals import post_save, post_delete
from .models import *
from django.dispatch import receiver
from PIL import Image
import os
from django.template.defaultfilters import slugify
from django.contrib.contenttypes.models import ContentType

def image_product(sender,instance,**kwargs):
    try:
        imagen = Image.open(instance.image.path)
        imagen.thumbnail((800,800))
        if imagen.format == 'PNG':
            bg = Image.new("RGB", (800,800), (255,255,255))
            bg.paste(imagen,imagen)
            bg.save(instance.image.path)
        else:
            imagen.save(instance.image.path)
    except Exception as e:
        pass

def del_image(sender,instance,**kwargs):
    os.remove(instance.image.path)

post_save.connect(image_product,sender=ProductImage)
post_delete.connect(del_image,sender=ProductImage)

def image_category(sender,instance,**kwargs):
    try:
        imagen = Image.open(instance.image.path)
        imagen.thumbnail((800,400))
        if imagen.format == 'PNG':
            bg = Image.new("RGB", (800,400), (255,255,255))
            bg.paste(imagen,imagen)
            bg.save(instance.image.path)
        else:
            imagen.save(instance.image.path)
    except Exception as e:
        pass

def del_category(sender,instance,**kwargs):
    os.remove(instance.image.path)

post_save.connect(image_category,sender=ProductCategory)
post_delete.connect(del_category,sender=ProductCategory)

def remove_index(sender,instance, **kwargs):
    from apps.search.models import Index
    ctype = ContentType.objects.get_for_model(instance)
    Index.objects.get(object_id=instance.id, content_type=ctype).delete()
post_delete.connect(remove_index,sender=Product)

def index_model(sender,instance,**kwargs):
    from apps.search.models import Index
    text =slugify(("%s-%s-%s")%(instance.name, instance.summary, instance.category.name))
    Index.objects.create(instance=instance, text=text)
post_save.connect(index_model,sender=Product)


import numpy as np
import torch

class ToTensor:
    """Convert ndarrays in sample to Tensors."""

    def __call__(self, sample):
        flow, rgb, y = sample['flow'], sample['rgb'], sample['y']
        # swap axis because
        # numpy image: N x H x W x C
        # torch image: N x C X H X W
        return {'flow': torch.from_numpy(np.asarray(flow).transpose((0, 3, 1, 2))),
                'rgb': torch.from_numpy(np.asarray(rgb).transpose((0, 3, 1, 2))),
                'y': torch.from_numpy(np.asarray(y))}

class RandomHorizontalFlip:
    """Randomly flips np arrays horizontaly"""

    def __init__(self, p=0.5):
        self.p = p

    def __call__(self, sample):
        flow, rgb, y = sample['flow'], sample['rgb'], sample['y']
        # torch images: N x C x H x W
        if np.random.rand() < self.p:
            flow = torch.flip(flow, dims=(-1,))
            rgb = torch.flip(rgb, dims=(-1,))

        return {'flow': flow, 'rgb': rgb, 'y': y}


class GameModule:
    def __init__(self, battle_code, game_controller):
        self.bc = battle_code
        self.gc = game_controller

        # map stuff
        self.map = self.gc.starting_map(self.gc.planet())

        # team stuff
        self.team_my = self.gc.team()
        self.team_enemy = self.bc.Team.Red if self.team_my == self.bc.Team.Blue else self.bc.Team.Blue

        if self.gc.planet() == self.bc.Planet.Earth:
            self.__init_my_and_enemy_pos_earth()

        # enemy detection
        self.__last_refresh = 0
        self.__get_refresh_units()

        print("[INFO]: Game Module init successful...")

    # Inverts location for enemy's perspective, assumes Earth and point symetry
    def __invert_earth_loc(self, loc):
        x = self.map.width - loc.x
        y = self.map.height - loc.y
        return self.bc.MapLocation(self.bc.Planet.Earth, x, y)

    @staticmethod
    def __loc_to_string(loc):
        return '(' + str(loc.x) + ',' + str(loc.y) + ')'

    # inits mine starting position and presumed enemy position
    def __init_my_and_enemy_pos_earth(self):
        # it's possible to have mutliple units at different places, but let's start with this one
        self.my_start = self.gc.my_units()[0].location.map_location()
        self.enemy_start = self.__invert_earth_loc(self.my_start)

        # TODO: Remove logs
        print('worker starts at ' + self.__loc_to_string(self.my_start))
        print('enemy worker presumably at ' + self.__loc_to_string(self.enemy_start))

    # internal-call method to get visible enemies. Should not be used more than once a round
    def __get_refresh_units(self):
        enemies = []
        my_units = []

        for unit in self.gc.units():
            if not unit.location.is_on_map():
                continue

            if unit.team == self.team_enemy:
                enemies.append(unit)
            else:
                my_units.append(unit)

        self.__last_enemies = enemies
        self.__last_my_units = my_units

    # === API ===
    def my_start_pos(self):
        # TODO: Solve Mars stuff
        return self.my_start

    def enemy_start_pos(self):
        # TODO: Solve Mars stuff
        return self.enemy_start

    def enemy_units(self):
        # Caching
        if self.__last_refresh != self.gc.round():
            self.__last_refresh = self.gc.round()
            self.__get_refresh_units()

        return self.__last_enemies

    def my_units(self):
        if self.__last_refresh != self.gc.round():
            self.__last_refresh = self.gc.round()
            self.__get_refresh_units()

        return self.__last_my_units

    ## Returns dictionary of leaderUnit: list of units in group
    def getGroups(self, effectiveDistance, selector=lambda u: true):
        groups = {}
        for leader in filter(selector, gc.my_units()):
            if not unit.selected:
                unit.selected = true;
                groups[leader] = gc.sense_nearby_units_by_team(unit.location, effectiveDistance, gc.team)

        return groups

    def countGroupless(self):
        counter = 0
        for unit in gc.my_units():
            if isInGroup(unit):
                counter += 1

        return counter

    def isInGroup(self, unit):
        for l, g in gc.groups:
            if unit in g:
                return True

        return False


"""
Module Functions:
    Plot King County House Rate data,
    in the form of sale listing cluster map and density heat map.
"""

import os
import webbrowser

import pandas as pd
import folium
from folium import plugins

# Set global settings and macros.
MAX_SHOW = 1000
HOUSE_URL = 'houses.html'
HOUSE_HEAT_URL = "househeatmap.html"
CRIME_HEAT_URL = "crimeheatmap.html"
pd.set_option('display.max_columns', None) # To display all columns

# Read in king county house data.
data = pd.read_csv("../../Data/kc_house_data.csv", parse_dates=['date'])
data['zipcode'] = data['zipcode'].astype(str)

## Create one map showing each listing of in house dataset and show in browser.

# Use folium Map function to plot underlying basic map.
house_map = folium.Map(location=[data['lat'].mean(), data['long'].mean()], zoom_start=10)
# Define clusters to show house clusters to add to the underlying houses_map.
marker_cluster = folium.MarkerCluster().add_to(house_map)
# Iteratively add interactive clusters to the basic map.
#	When mouse-over the cluster, show house listing information:
#	sqft, price.
for iters, row in data[0:MAX_SHOW].iterrows():
    folium.Marker([row["lat"], row["long"]],
                  popup="{0} sqft:  Sold for $ {1}"\
                  .format(row["sqft_living"], row["price"])).add_to(marker_cluster)

# Save the house cluster map to a html and open it in broswer.
house_map.save(HOUSE_URL)
webbrowser.open('file://' + os.path.realpath(HOUSE_URL), new=2)

## Create one map showing the frequency of house sales and show in browser.

# Use folium Map function to plot underlying basic map
houses_heatmap = folium.Map(location=[data['lat'].mean(), data['long'].mean()], zoom_start=10)
# Add heatmap on top of the basic map.
houses_heatmap.add_children(
    plugins.HeatMap([[row["lat"], row["long"], row["price"]]
                     for iters, row in data.iterrows()],
                    min_opacity=0.5, max_zoom=18, radius=8)) #[0:MAX_SHOW]

# Save the house sale frequency heat map to a html and open it in broswer.
houses_heatmap.save(HOUSE_HEAT_URL)
webbrowser.open('file://' + os.path.realpath(HOUSE_HEAT_URL), new=2)


#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Ringtail script for updating v1.0.0 databases to v1.1.0
#

import argparse
from ringtail import RingtailCore
import logging
import sys

if __name__ == "__main__":
    logging.basicConfig(
        level=logging.INFO, stream=sys.stdout, filemode="w", format="%(message)s"
    )
    # get name(s) of dbs to update from command line
    parser = argparse.ArgumentParser(
                    prog='rt_db_v100_to_v110',
                    description='Given one or multiple Ringtail databases made with v1.0.0, will update them to be compatible with v1.1.0')

    parser.add_argument('-d', '--database', help="Database file(s) made with Ringtail v1.0.0 to update to v1.1.0", nargs="+", type=str, action="store")
    args = parser.parse_args()

    consent = False

    for db in args.database:
        defaults = RingtailCore.get_defaults()
        defaults["storage_opts"]["values"]["db_file"] = db

        with RingtailCore("sqlite", defaults) as rtcore:
            consent = rtcore.update_database(consent)


# -*- coding: utf-8 -*-
# Time : 2020/1/3 0003  8:59
# Author : dengfan
import json, os, cv2
import numpy as np
import copy

from PIL import Image, ImageDraw, ImageFont

'''
 index = ["PGBX", "PGDD", "PMZC", "PGHB", "PGPS", "PMYC", "PGDX", "BTQZ", "BTQP", "BTWX" ]
'''

classes = {
    0: '背景',
	1: 'PGPS',#'瓶盖破损',
	2: 'PGBX',#'瓶盖变形',
	3: 'PGHB',#'瓶盖坏边',
	4: 'PGDX',#'瓶盖打旋',
	5: 'PGDD',#'瓶盖断点',
	6: 'BTWX',#'标贴歪斜',
	7: 'BTQZ',#'标贴起皱',
	8: 'BTQP',#'标贴气泡',
	9: 'PMZC',#'喷码正常',
	10: 'PMYC'#'喷码异常'}
}

def read_json(files_path):
    with open(files_path, encoding='utf-8') as f:
        dic = json.load(f)
        # dic = json.loads(setting)
    return dic

def get_dic(data):
    data = read_json(path)
    for i in data:
        print(i)

    images = data['images']
    data_images = {}
    for dic in images:
        file_name = dic['file_name']
        id = dic['id']
        if id in data_images:
            print(id)
            print(file_name)
            raise Exception('文件{} 已经存在，出现重复编码！！！'.format(file_name))
        data_images[id] = file_name

    annotations = data['annotations']
    out = {}
    for ann in annotations:
        name = data_images[ann['image_id']]
        bbox = ann['bbox']
        category = ann['category_id']
        score = ann['score']
        if name not in out:
            out[name] = [[bbox[0], bbox[1], bbox[2], bbox[3], category, score], ]
        else:
            out[name].append([bbox[0], bbox[1], bbox[2], bbox[3], category, score])
    return out

if __name__ == "__main__":
    path = r'../process_after/result.json'
    img_path = r'E:\study_code\TC_CQ\data\test\images'
    save_path = r'./show_imgs2'

    data = get_dic(read_json(path))

    num = 0
    for name in data:
        num += 1
        # if num > 10:
        #     break
        name_path = os.path.join(img_path, name)
        print(name_path)
        bboxes = data[name]
        img = cv2.imread(name_path)
        for bbox in bboxes:
            if bbox[5] < 0.15:
                continue
            img = cv2.rectangle(img, (int(bbox[0]), int(bbox[1])),(int(bbox[0]+bbox[2]), int(bbox[1]+bbox[3])),
                                (0, 0, 255), thickness=1)

            img = cv2.putText(img, classes[bbox[4]]+'-'+ str(bbox[5])[:6], (int(bbox[0]), int(bbox[1]) -5),
                              cv2.FONT_HERSHEY_COMPLEX, 0.5, (0, 255, 0), 1)
        save_name = os.path.join(save_path, name)
        cv2.imwrite(save_name, img)


from mdp_comm.bluetooth.bluetooth_server import BluetoothServer
from mdp_comm.image.image_server import ImageServer
from mdp_comm.robot.stm32 import STM32
# from Pathfinding.main import pathfind
from pathing import pathfind

from rpimessageparserV4 import androidToRpi
from obstacle_storage import ObstacleStorage

import time

path_map = {
    'RIGHT_FWD': [("r", 90)],
    'LEFT_FWD': [("l", 90)],
    '3PT_LEFT': [("l",)],
    '3PT_TURN_AROUND':[("l", ), ("l", )],
    '3PT_RIGHT': [("r",)],
    'REVERSE': [("b",)],
    'LEFT_RVR': [("bl",)],
    'RIGHT_RVR': [("br",)],
    'FORWARD': [("f", )],
    'f':[('f', )],
    'r':[('r', )],
    'b':[('b', )],
    'l':[('l', )],

}

def path_map_get(key):
    if key not in path_map:
        val = [key]
    else:
        val = path_map[key]

    print(val)
    return val

facing_to_dir = {
    (0, 1): "N",
    (0, -1): "S",
    (1, 0): "E",
    (-1, 0): "W",
}

# N E S W
right_update = {
    (0, 1): (1, 0),
    (1, 0): (0, -1),
    (0, -1): (-1, 0),
    (-1, 0): (0, 1),
}

# N W S E
left_update = {
    (0, 1): (-1, 0),
    (-1, 0): (0, -1),
    (0, -1): (1, 0),
    (1, 0): (0, 1),
}

class Application:
    def __init__(self) -> None:
        self.image_server = ImageServer(self.image_server_callback)
        self.bluetooth_server = BluetoothServer(self.bluetooth_server_callback)
        self.robot = STM32()
        self.obstacle_storage = ObstacleStorage()

        self.robot.connect()
        self.image_server.run()
        self.bluetooth_server.run()
        while not self.image_server.is_connected() and self.bluetooth_server.is_connected():
            time.sleep(1)
            print("waiting for connection")

        self.pos = (1, 1)
        self.facing = (0, 1)
        self.obstacle_num = 1

    def update_location(self, move):

        x, y = self.pos
        d_x, d_y = self.facing

        assert isinstance(move, tuple)

        if move[0] == "r":
            self.facing = right_update[self.facing]

        elif move[0] == "l":
            self.facing = left_update[self.facing]

        elif move[0] == "f":
            if len(move) > 1:
                self.pos = (x+ (move[1]/10)*d_x, y+ (move[1]/10)*d_y)
            else:
                self.pos = (x+ d_x, y+ d_y)

        elif move[0] == "b":
            if len(move) > 1:
                self.pos = (x- (move[1]/10)*d_x, y- (move[1]/10)*d_y)
            else:
                self.pos = (x- d_x, y- d_y)

    def image_server_callback(self, msg):
        print(msg)
        msg = msg.decode()
        if msg == "None":
            pass
        else:
            print(f"TARGET,<{self.obstacle_num}>,<{msg}>")
            self.bluetooth_server.send(f"TARGET,<{self.obstacle_num}>,<{msg}>")
            self.obstacle_num += 1

    def bluetooth_server_callback(self, msg):

        msg = msg.decode()

        parsed = androidToRpi(msg)
        if parsed is None:
            return
        elif parsed[0] == "moving":
            self.send_robot_command(parsed[1])
        elif parsed[0] == "Obstacle":
            _, x, y, obstacle_id, obstacle_dir = parsed
            self.obstacle_storage.update_obstacle(obstacle_id, obstacle_dir, int(x), int(y))
            print(self.obstacle_storage._obstacles)
        elif parsed[0] == "starting":
            pass
        elif parsed[0] == "discover":
            self.start_discover()
        elif parsed[0] == "fastest":
            self.start_fastest()
        elif parsed[0] == "shutdown":
            import subprocess
            print("shutting down")
            subprocess.run(["shutdown", "0"])
        elif parsed[0] == "clear":
            self.obstacle_storage = ObstacleStorage()

    def send_robot_command(self, direction, value=None):
        if direction == "l":
            if value is None:
                self.robot.write("\rL\r")
            else:
                self.robot.write(f"\rL{value}\r")

        elif direction == "r":
            if value is None:
                self.robot.write("\rR\r")
            else:
                self.robot.write(f"\rR{value}\r")

        elif direction == "f":
            if value is None:
                value = 10
            self.robot.write(f"\rF{value}\r")

        elif direction == "b":
            if value is None:
                value = 10
            self.robot.write(f"\rB{value}\r")

        elif direction == "br":
            self.robot.write(f"\rA90\r")

        elif direction == "bl":
            self.robot.write(f"\rD90\r")
        elif direction == "fr":
            self.robot.write(f"\rR90\r")

        elif direction == "fl":
            self.robot.write(f"\rL90\r")

        if self.robot.read() != "C":
            print("Panic! expected C")
        else:
            print("Completed move...")

    def _move_path(self, path, capture=True):
        for move in path:
            self.send_robot_command(*move)
            if capture:
                self.image_server.capture_and_send()

            self.update_location(move)
            print(f"ROBOT,<{int(self.pos[0])}>,<{int(self.pos[1])}>,<{facing_to_dir[self.facing]}>".encode())
            self.bluetooth_server.send(f"ROBOT,<{int(self.pos[0])}>,<{int(self.pos[1])}>,<{facing_to_dir[self.facing]}>".encode())

    def start_fastest(self):
        pass

    def start_discover(self):
        self.pos = (1, 1)
        self.facing = (0, 1)
        self.obstacle_num = 1

        data = self.obstacle_storage.extract_required_format()
        result = pathfind(*data)

        # path = result['pathfinding']['moves']
        path = result
        for sub_path in path:
            print(sub_path)
            sub_path = self._parse_path(sub_path)
            self._move_path(sub_path, capture=False)
            self.image_server.capture_and_send()

    def _parse_path(self, path):
        pth = []

        for p in path:

            # pth += path_map[p]
            pth += path_map_get(p)

        print(pth)
        return pth

if __name__ == "__main__":
    app = Application()


from __future__ import division
import re
import os

from collections import defaultdict
from heapq import nlargest

def ScanFiles(path): #fun
	for dirs,subdir,files in os.walk(path):
		for fn in files:
			if fn.find("Packet.") >= 0:
				yield fn

def PrintToppestInfo(dict, topIndex, OutPutFile, printName):
	Toppest_List = nlargest(topIndex, dict.items(),key=lambda x: x[1]);
	print("Toppest_"+ printName +": ")
	OutPutFile.write("Toppest_"+printName+": \n");
	OutPutFile.write("rank packetId "+printName+" PacketName \n");
	index = 1;
	for item in Toppest_List:
			print(index, item[0],item[1])
			OutPutFile.write(str(index)+" "+str(item[0])+" "+str(item[1])+" "+ pakcetNameDict[item[0]] + "\n");
			index += 1;
	OutPutFile.write("\n");
	return;

receivecountDict = defaultdict(int);
receivesizeDict = defaultdict(int);
sendcountDict = defaultdict(int);
sendsizeDict = defaultdict(int);
pakcetNameDict = defaultdict(str);
packetSizeDict = defaultdict(int);

if not os.path.isdir("Result"):
	os.makedirs("Result");
OutPutFile = open("Result\PacketLogAnalyseResult.txt", "w");

TopIndex = 10;
print(os.getcwd())
with open(os.getcwd()+"\Config\Config.txt") as configFile:
	for configline in configFile:
			TopIndex = (int)(configline);
			print(TopIndex);

os.chdir("Log");
allFiles = list(ScanFiles(os.getcwd()))
for lfile in allFiles:
	with open(lfile) as logfile:
		for line in logfile:
			dotPos = line.find(",",0);
			if dotPos > 0 and line[0:dotPos].isdigit():
				packetId, szreceivecount, szreceivesize, szsendcount, szsendsize, PacketName, noUse = line.split(",",6);
				receivecount = (int)(szreceivecount);
				receivesize = (int)(szreceivesize);
				sendcount = (int)(szsendcount);
				sendsize = (int)(szsendsize);
				receivecountDict[packetId] += receivecount;
				receivesizeDict[packetId] += receivesize;
				sendcountDict[packetId] += sendcount;
				sendsizeDict[packetId] += sendsize;
				pakcetNameDict[packetId] = PacketName;
				if receivecount > 0 and receivesize >0:
					packetSizeDict[packetId] = receivesize/receivecount;
				elif sendcount >0 and sendsize >0:
					packetSizeDict[packetId] = sendsize/sendcount;

OutPutFile.write("Toppest "+str(TopIndex)+" OutPut: \n");

PrintToppestInfo(receivecountDict,TopIndex, OutPutFile, "receivecount");

PrintToppestInfo(receivesizeDict,TopIndex, OutPutFile, "receivesize");

PrintToppestInfo(sendcountDict,TopIndex, OutPutFile, "sendcount");

PrintToppestInfo(sendsizeDict,TopIndex, OutPutFile, "sendsize");

PrintToppestInfo(packetSizeDict,TopIndex, OutPutFile, "packetsize");

OutPutFile.close();


""" Search flat results on wg-gesucht.de """
from __future__ import unicode_literals, absolute_import, generators, print_function

from lxml import html
from urllib import urlencode
import requests
import time
import random
import locale

locale.setlocale(locale.LC_TIME, str('de_DE.UTF-8'))

class WgGesucht(object):
    """ WG-Gesucht Screen Scraper. """
    def __init__(self):
        """ Initialize with base url """
        self.base_url = 'http://www.wg-gesucht.de/wohnungen-in-Berlin.8.2.0.{}.html?'

    def search(self, category, rent_type, minSize, maxPrice, minRooms, maxRooms, exc, balcony, pets, furnished):
        """ Search using a get request including flat details.
        :param category: type of flat
        :param rent_type: type of contract
        :param minSize: minimum size in square meters
        :param maxPrice: max rent price
        :param minRooms: minimum number of rooms
        :param maxRooms: max number of rooms
        :param exc: exchange flat y/n (optional)
        :param balcony: balcony y/n (optional)
        :param pets: pets allowed (optional)
        :param furnished: furnished y/n (optional)
        """
        params = {'offer_filter': 1, 'city_id': 8,
                  'category': category, 'rent_type': rent_type,
                  'sMin': minSize, 'rMax': maxPrice,
                  'rmMin': minRooms, 'rmMax': maxRooms,
                  'exc': exc, 'bal': balcony,
                 'pet': pets, 'fur': furnished}

        session = requests.Session()
        response = session.get('{}{}'.format(self.base_url.format('0'), urlencode(params)))
        print(response.url)
        results = self.parse_response(response)[0]

        max_page = self.has_more_pages(response)
        if max_page:
            stop = False
            for i in range(1, max_page):
                time.sleep(random.randint(1, 3))
                if stop:
                    break
                r = session.get(self.base_url.format(i))
                print(r.url)
                next_pages, stop = self.parse_response(r)
                results.extend(next_pages)
        return results

    def has_more_pages(self, response):
        page = html.fromstring(response.content)
        try:
            max_page = int(page.xpath('//a[@class="a-pagination"][last()]')[-1].text.replace('\n', '').strip())
        except Exception as IndexError:
            print('Only one page results')
            max_page = False
        return max_page

    def grab_xpath_text(self, element, xpath):
        """ Given an element and xpath pattern, return text content.
        :param element: lxml element
        :param xpath: string
        returns string
        """
        data = element.xpath(xpath)
        if len(data) == 1:
            return data[0].text
        elif len(data) > 1:
            return [x.text for x in data]
        return ''

    def parse_response(self, response):
        """ Given a requests response object, return a list of dictionaries
        containing the pertinent flat info.
        :params response: response obj
        returns list of dictionaries
        """
        stop = False
        page = html.fromstring(response.content)
        results = page.xpath('//table/tbody/tr')
        active_inactive = [self.grab_xpath_text
                           (res, 'td[contains(@class, "datum")]/a/span').replace('\n', '').strip() not in 'inaktiv' for res in results]
        final_results = []
        for is_active, res in zip(active_inactive[2:], results[2:]):
            if not is_active:
                stop = True
                break
            item_dict = {}
            item_dict['rooms'] = self.grab_xpath_text(
                res, 'td[contains(@class, "zimmer")]/a/span').replace('\n', '').strip()
            if not item_dict['rooms']:
                item_dict['rooms'] = '1'
            item_dict['Free from'] = self.grab_xpath_text(
                res, 'td[contains(@class, "freiab")]/a/span').replace('\n', '').strip()
            item_dict['Free until'] = self.grab_xpath_text(
                res, 'td[contains(@class, "freibis")]/a/span').replace('\n', '').strip()
            item_dict['Rent price'] = self.grab_xpath_text(
                res, 'td[contains(@class, "miete")]/a/span/b').replace('\n', '').strip()
            item_dict['Size'] = self.grab_xpath_text(
                res, 'td[contains(@class, "groesse")]/a/span').replace('\n', '').strip()
            item_dict['District'] = self.grab_xpath_text(
                res, 'td[contains(@class, "stadt")]/a/span').replace('\n', '').strip()
            item_dict['Link'] = 'http://www.wg-gesucht.de/'+res.get('adid')
            final_results.append(item_dict)
        return final_results, stop


from tkinter import *
from tkinter import ttk
from tkinter import filedialog
import os
import socket

listPlace = 0

def get_folder_path():
    folder_selected = filedialog.askdirectory()
    if folder_selected == '':
        get_folder_path()
    else:
        folderPath.set(folder_selected)
    # print(folderPath.get().split('/')[-1])

def go_back():
    gui.destroy()
    gui.quit()
    os.system('python main.py')

def start_key():

    print("Getting File")
    os.system(f'py receiver_sockets.py --port {int(portText.get())} --file {folderPath.get()} --host {default_ip_text.get()}')

def just_started():
    global listPlace
    listPlace += 1
    text_area.insert(listPlace, listPlace)
    text_area.select_clear(text_area.size() - 2)
    text_area.select_set(END)
    text_area.yview(END)
    start_key()

def disable_port():
    buttonRandom.config(state = DISABLED)
    portEntry.config(state = DISABLED)

gui = Tk()
gui.geometry("300x260")
gui.iconbitmap('logo.ico')

gui.title("Receive Files")
default_ip_text = StringVar()
hostname = socket.gethostname()
local_ip = socket.gethostbyname(hostname)
default_ip_text.set(local_ip)

folderPath = StringVar()
folderPath.set('Folder Name Here')

portText=StringVar()

c = ttk.Button(gui, text="Back", command=go_back)
c.place(x=0, y=5)

a = Label(gui, text="File")
a.place(x=0, y=35)
btnFind = ttk.Button(gui, text="Select Folder", command=get_folder_path)
btnFind.place(x=50, y=35)
E = Entry(gui, textvariable=folderPath, state='disabled')
E.place(x=150, y=35)

portLbl = Label(gui, text="Port")
portLbl.place(x=0, y=65)
buttonRandom = ttk.Button(gui, text="Set Port", command=disable_port)
buttonRandom.place(x=50, y=65)
portEntry = Entry(gui, textvariable=portText)
portEntry.place(x=150, y=65)

keyLbl = Label(gui, text="IP")
keyLbl.place(x=0, y=95)
KeyEntry = Entry(gui, textvariable=default_ip_text,width=37)
KeyEntry.place(x=50, y=95)

startButton = ttk.Button(gui, text="Start", command=just_started)
startButton.place(x=100, y=125)

progress = ttk.Progressbar(gui, orient=HORIZONTAL, value=10, length=290, mode='determinate')
progress.place(x=0, y=160)

log = Label(gui, text="Logs", font=("Times New Roman", 7))
log.place(x=0, y=185)
text_area = Listbox(gui, width=48, height=3, font=("Times New Roman", 10))
text_area.place(x=0, y=205)
text_area.insert(listPlace, "Press Start To Start Program")
listPlace += 1
text_area.insert(listPlace, "When u press the start from key this program will receive ")
listPlace += 1
text_area.insert(listPlace, 'all data from other computer sending with the same port')
text_area.select_set(END)
text_area.yview(END)

get_folder_path()
gui.resizable(False, False)
gui.mainloop()


#from leetcode medium difficulty https://leetcode.com/problems/remove-nth-node-from-end-of-list/

# Definition for singly-linked list.
# class ListNode:
#     def __init__(self, val=0, next=None):
#         self.val = val
#         self.next = next
class Solution:
    def removeNthFromEnd(self, head: ListNode, n: int) -> ListNode:

        comp_arr = []
        current = head

        while current is not None:

            comp_arr.append(current.val)
            current = current.next

        current = head
        prev = None
        incrementor = 0
        while current is not None:
            if head.next is None and n == 1:
                head = None
                return head

            elif head.next is None and n == 0:
                return head

            elif n == len(comp_arr):
                head = head.next
                del current
                return head
            else:

                prev = current
                current = current.next
                incrementor += 1

                if current.val == comp_arr[len(comp_arr)-n] and incrementor == len(comp_arr) - n:
                    print(f'in: {comp_arr}')
                    print(f'prev: {prev} current: {current}')
                    prev.next = current.next
                    # current.next = None???
                    del current
                    return head


import numpy as np
import torch

from escnn import nn, gspaces, group

from .mlp import SONMLPTensor

class ImplicitKernelSON(torch.nn.Module):
    """
    Steerable kernel:
        3D coordinates -> harmonic polynomicals -> + edge attributes -> implicit kernel.
    """
    def __init__(
        self,
        in_repr: nn.FieldType,
        out_repr: nn.FieldType,
        edge_repr: nn.FieldType,
        hp_order: int = None,
        edge_distr: list = [None, None],
        **mlp_kwargs):
        """
        Args:
            in_repr (nn.FieldType): representation of the convolution input.
            out_repr (nn.FieldType): representation of the convolution output.
            edge_repr (nn.FieldType): representation of the edge attributes.
            hp_order (int): order of harmonic polynomials.
            mlp_kwargs (dict): keyword arguments for the MLP.
        """
        super().__init__()
        if hp_order is None:
            # if the order of HPs is not given, we take the maximum frequency of the input representation
            ik_irreps = out_repr.representations[0].tensor(in_repr.representations[0]).irreps # ψ1 ⊗ ψ2
            hp_order = max(ik_irreps, key=lambda x: x[1])
        self.edge_distr = edge_distr
        self.G = group.o3_group()
        self.gspace = gspaces.no_base_space(self.G)
        # input: 3D coordinates + additional edge attributes
        self.pos_type = self.gspace.type(self.G.standard_representation())
        self.edge_repr = edge_repr
        self.hp = nn.HarmonicPolynomialR3(hp_order, 'o3')
        # restrict output of HPs to a subgroup defined by the input representation
        self.subgroup_id = in_repr.gspace._sg_id[:-1]
        # harmonic polynomial output + edge attributes
        mlp_in_repr = [i for i in self.hp.out_type.representations] + edge_repr
        # we restrict the representation of input type and add attribute representations if needed
        mlp_in_type = self.gspace.type(*mlp_in_repr).restrict(self.subgroup_id)
        # batch normalization for input coordinates
        self.bn_coords = nn.IIDBatchNorm1d(self.pos_type)
        # irrep-wise batch normalization for HPs output
        decomposed_type = self.gspace.type(*[self.hp.out_type.fibergroup.irrep(*id) for id in self.hp.out_type.irreps])
        self.bn_hp = nn.IIDBatchNorm1d(decomposed_type)
        # function that applies batch normalization to each irrep independently
        self.scale_hp = lambda x: self.hp.out_type(self.bn_hp(decomposed_type(x.tensor)).tensor)
        self.mlp = SONMLPTensor(in_type=mlp_in_type,
                                in_repr=in_repr,
                                out_repr=out_repr,
                                initialize=True,
                                **mlp_kwargs)
        self.c_in, self.c_out  = len(in_repr), len(out_repr)
        self.delta_in, self.delta_out = int(in_repr.size / len(in_repr)), int(out_repr.size / len(out_repr))
        # learn scale of the kernel
        self.inv_scale = torch.nn.Parameter(torch.tensor(1.0))
        # kernel's output initial standard deviation (for generalized He initialization)
        kstd = 1 / self.kernel_std(1000)
        self.register_buffer("inv_std", kstd)
        # scale parameter
        self.init_factor = 2 / np.sqrt(self.c_in * self.c_out)

    def forward(self, coords: torch.Tensor, edge_attr: torch.Tensor, init: bool = True):
        """
        Forward pass of the kernel:
            Input -> Scale coords -> HPs -> MLP -> He init -> filter.
        Args:
            coords (torch.Tensor): 3D coordinates of the input points.
            edge_attr (torch.Tensor): edge attributes.
            init (bool): whether to re-scale the filter (corresponds to He initialization).
        """
        x = self.transform_coords(coords)
        if edge_attr is not None:
            x = nn.GeometricTensor(torch.cat([x.tensor, edge_attr], 1), self.mlp.in_type)
        else:
            x = nn.GeometricTensor(x.tensor, self.mlp.in_type)
        # compute c_in * delta_in * c_out * delta_out filter based on input features
        x = self.mlp(x)
        x = self.transform_mlp(x, coords, init)
        # reshape the vector to d_out x d_in filter at last
        return x.reshape((-1, self.c_out * self.delta_out, self.c_in * self.delta_in))

    def transform_coords(self, coords: torch.Tensor):
        """
        Transform coordinates of the input points:
            BatchNorm -> Harmonic polynomials -> scaling.
        """
        coords = nn.GeometricTensor(coords, self.pos_type)
        coords = self.bn_coords(coords) # normalize input coordinates
        x = self.compute_hp(coords)
        x = self.scale_hp(x)
        return x

    def transform_mlp(self, x: nn.GeometricTensor, coords: torch.Tensor, init: bool):
        """
        Transform the output of MLP:
            Scale MLP's output -> reshape the output to c_out x delta_out x c_in x delta_in filter -> He init.
        Args:
            x (nn.GeometricTensor): output of the MLP.
            coords (torch.Tensor): 3D coordinates of the input points.
            init (bool): whether to re-scale the filter (corresponds to He initialization).
        """
        # multiply with exp radial function
        x = x.tensor
        x = self.scale_mlp(x, coords)
        x = x.view((-1, self.c_out, self.c_in, self.delta_out, self.delta_in))
        # normalize the filter to yield He weight initialization
        x = self.he_initialize(x) if init else x
        return x.permute((0, 1, 3, 2, 4))

    def scale_mlp(self, x, coords):
        """
        Scale MLP's output according to the coordinates of the input point.
        Similar to radial function of SO(3) steerable kernels: https://arxiv.org/abs/1807.02547.
        """
        squared_norms = torch.sum(coords**2, dim=1, keepdim=True)
        # Compute the exponential factor of the operation
        exponential_factor = torch.exp(-0.5 * squared_norms * self.inv_scale**2)
        # Apply the element-wise multiplication of `x` and the exponential factor
        return x * exponential_factor

    def compute_hp(self, x: nn.GeometricTensor):
        """
        Compute scaled harmonic polynomials that are restricted to the subgroup of the input representation.
        """
        return self.hp(x).restrict(self.subgroup_id)

    def he_initialize(self, x: torch.Tensor):
        """
        Generalized He weight initialization of filters calculated by the implicit kernel.
        The initialization condition is taken from https://arxiv.org/pdf/1711.07289.pdf (see Appendix B).
        For the forward pass: Var[w] = 2 / (C * Q) where
            C is the number of input channels,
            Q is the number of basis functions,
            filters are normalized to unit norm.
        """
        return x * self.inv_std * self.init_factor

    def sample_kernel(self, n_points: int):
        """
        Sample points in space randomly and compute the value of kernel at each point.
        Draws samples from a normal distribution with mean 0 and variance 1.
        """
        device = next(self.parameters()).device
        with torch.no_grad():
            # randomly generated coordinates in 3D
            if self.edge_distr[0] is None:
                coords = torch.randn(n_points, 3).to(device)
            else:
                coords = self.edge_distr[0].sample([n_points]).to(device)

            if self.edge_distr[1] is None:
                edge_attr = torch.randn(n_points, sum(reps.size for reps in self.edge_repr)).to(device)
            else:
                edge_attr = self.edge_distr[1].sample([n_points]).to(device)

            return self(coords, edge_attr, init=False)

    def kernel_std(self, n_points: int):
        """
        Computes channel-wise variance for the implicit kernel;
        asserts that there are no zero variance channels.
        """
        x = self.sample_kernel(n_points)
        x = x.reshape((-1, self.c_out, self.delta_out, self.c_in, self.delta_in))
        return x.std(dim=(0,2,4)).reshape(1, self.c_out, self.c_in, 1, 1)


import conftest  # Add root path to sys.path
import os
import matplotlib.pyplot as plt
from PathPlanning.SpiralSpanningTreeCPP \
    import spiral_spanning_tree_coverage_path_planner

spiral_spanning_tree_coverage_path_planner.do_animation = True

def spiral_stc_cpp(img, start):
    num_free = 0
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            num_free += img[i][j]

    STC_planner = spiral_spanning_tree_coverage_path_planner.\
        SpiralSpanningTreeCoveragePlanner(img)

    edge, route, path = STC_planner.plan(start)

    covered_nodes = set()
    for p, q in edge:
        covered_nodes.add(p)
        covered_nodes.add(q)

    # assert complete coverage
    assert len(covered_nodes) == num_free / 4

def test_spiral_stc_cpp_1():
    img_dir = os.path.dirname(
        os.path.abspath(__file__)) + \
        "/../PathPlanning/SpiralSpanningTreeCPP"
    img = plt.imread(os.path.join(img_dir, 'map', 'test.png'))
    start = (0, 0)
    spiral_stc_cpp(img, start)

def test_spiral_stc_cpp_2():
    img_dir = os.path.dirname(
        os.path.abspath(__file__)) + \
        "/../PathPlanning/SpiralSpanningTreeCPP"
    img = plt.imread(os.path.join(img_dir, 'map', 'test_2.png'))
    start = (10, 0)
    spiral_stc_cpp(img, start)

def test_spiral_stc_cpp_3():
    img_dir = os.path.dirname(
        os.path.abspath(__file__)) + \
        "/../PathPlanning/SpiralSpanningTreeCPP"
    img = plt.imread(os.path.join(img_dir, 'map', 'test_3.png'))
    start = (0, 0)
    spiral_stc_cpp(img, start)

if __name__ == '__main__':
    conftest.run_this_test(__file__)


# -*- coding: utf-8 -*-
# Computing Persistent Homology and its histogram
## install cripser by
## pip install git+https://github.com/shizuo-kaji/CubicalRipser_3dim

import os,glob
import numpy as np
from tqdm import tqdm
import seaborn as sns
import matplotlib.pyplot as plt
from sklearn.neighbors import KernelDensity
from scipy.stats import gaussian_kde
import argparse,json
from scipy.ndimage.morphology import distance_transform_edt
from skimage import feature,morphology
from skimage.filters import threshold_otsu, threshold_niblack,threshold_sauvola, scharr
from PIL import Image
from multiprocessing import Pool
from functools import partial
from arguments import arguments
try:
    import persim
    import cripser
    from gudhi.representations import Landscape
except:
    pass

# preprocess image before computing PH
def preprocess_image(img, gradient=False, img_size=None, filtration=None, origin=(0,0), binarisation='otsu'):
    im = img

    if img_size:
        import skimage.transform
        im = skimage.transform.resize(im,(img_size,img_size))

    if gradient:
        import cv2
        #im = feature.canny(im, sigma=10)
        #im = scharr(im)
        if len(im.shape)==2: # temporaly convert to BGR to use cv2 functions
            im = cv2.cvtColor(im,cv2.COLOR_GRAY2BGR)
        denoised = cv2.fastNlMeansDenoisingColored(im, None, 30, 10, 7, 21)
        im = cv2.cvtColor(denoised, cv2.COLOR_BGR2GRAY)
        v = np.median(im)
        sigma = 0.9 # 0.33
        lower = int(max(0, (1.0 - sigma) * v))
        upper = int(min(255, (1.0 + sigma) * v))
        im = cv2.Canny(im, lower, upper)

    if len(im.shape)>2:
        im = np.dot(im[...,:3], [0.2989, 0.5870, 0.1140])

    if filtration is not None:
        if im.max()==im.min():
            return(np.zeros_like(im))
        if gradient:
            bw_img = im
        else:
            if binarisation=='sauvola':
                bw_img = (im >=threshold_sauvola(im))
            elif binarisation=='niblack':
                bw_img = (im >= threshold_niblack(im))
            else:
                bw_img = (im >= threshold_otsu(im))
        #bw_img = morphology.area_opening(bw_img, area_threshold=8, connectivity=2)

        if filtration=='binarise':
            return(bw_img)
        elif filtration == 'distance':
            dt_img = distance_transform_edt(~bw_img) # distance from the foreground 0
            #print(dt_img.max())
        elif filtration == 'signed_distance':
            dt_img = distance_transform_edt(~bw_img)-distance_transform_edt(bw_img)
        elif filtration in ['downward','upward']:
            null_idx = (bw_img == 0)
            ## height transform
            if len(im.shape) == 3: #(z,y,x)
                h = np.arange(im.shape[0]).reshape(-1,1,1)
            else:
                h = np.arange(im.shape[0]).reshape(-1,1)
            if filtration=='upward':
                h = np.max(h) - h
            dt_img = (bw_img * h)
            dt_img[null_idx] = np.max(h)
        elif 'radial' in filtration:
            null_idx = (bw_img == 0)
            h = np.linalg.norm(np.stack(np.meshgrid(*map(range,im.shape),indexing='ij'),axis=-1)-np.array(origin), axis=-1)
            dt_img = (bw_img * h)
            if filtration=='radial_inv':
                dt_img = np.max(dt_img) - dt_img
            else:
                dt_img[null_idx] = np.max(h)

        dt_img *= 256/dt_img.shape[0]  # scaling normalisation
        return(dt_img)
    else:
        return(im)

## computing PH of an image
def comp_PH(img, gradient=True, img_size=None, filtration=None):
    im = preprocess_image(img, gradient=gradient, img_size=img_size, filtration=filtration)
    pd = cripser.computePH(im.astype(np.float64))
    #pd = cripser.computePH(im,maxdim=args.maxdim,top_dim=args.top_dim,embedded=args.embedded)
    #print(sum(pd[:,0] == 1))
    #print(im.shape)
    return(pd)

def comp_save_PH(fname, args):
    bfn = os.path.splitext(os.path.basename(fname))[0]
    img = np.array(Image.open(fname).convert('L'),dtype=np.float64)
    ph = comp_PH(img, gradient=args.gradient, filtration=args.filtration, img_size=args.img_size)
    np.save(os.path.join(args.output, bfn), ph)

def comp_landscape(ph,  dim, min_birth=None, max_birth=None, max_life=None,n=5):
    res = []
    for d in [0,1]:
        pds = ph[ph[:,0] == d, 1:3]
        #pds[:,1] = pds[:,0]+(np.clip(pds[:,1]-pds[:,0],0,max_life))
        res.append(Landscape(num_landscapes=n, resolution=dim[d]//n, sample_range=[min_birth[d],max_birth[d]]).fit_transform([pds]).ravel().astype(np.float32))
    return(np.sqrt(np.concatenate(res)))

def comp_persistence_image(ph, args=None):
    pims = []
    for d in [0,1]:
        s = np.sqrt((args.max_birth[d]-args.min_birth[d])*args.max_life[d]/args.num_bins[d])
        p = int((args.max_birth[d]-args.min_birth[d])/s)
        q = int(args.max_life[d]/s)
        while p*q < args.num_bins[d]:
            s = max((args.max_birth[d]-args.min_birth[d])/(p+1), args.max_life[d]/(q+1))
            p = int((args.max_birth[d]-args.min_birth[d])/s)
            q = int(args.max_life[d]/s)
        pim = persim.PersistenceImager(birth_range=(args.min_birth[d],args.max_birth[d]),
            pers_range=(0,args.max_life[d]),pixel_size=s,
            kernel_params={'sigma': [[args.persImg_sigma, 0.0], [0.0, args.persImg_sigma]]},
            weight_params={'n': args.persImg_weight})
        p = (ph[ph[:,0]==d])[:,1:3] # extract dim=d cycles
        life = p[:,1]-p[:,0]
        life = np.clip(life,a_min=None,a_max=args.max_life[d])
        p[:,1] = life
        pi = pim.transform(p, skew=False)
        #print(s, args.num_bins[d])
        #print(pi.shape)
        pi = pi.ravel()
        #pi = np.pad(pi,(0,args.num_bins[d]))
        pi = pi[:args.num_bins[d]]
        pi = np.abs(pi) ** args.persImg_power # to suppress overflow during learning
        pims.append(pi.astype(np.float32))
    return(pims)

def comp_betticurve(ph, dim, min_birth=None, max_birth=None, max_life=None):
    res = []
    for d in range(2):
        pds = ph[ph[:,0] == d, 1:3]
        mlife = (np.clip(pds[:,1]-pds[:,0],0,max_life[d]))

        res.append(np.zeros(dim[d]))
        for i,th in enumerate(np.linspace(min_birth[d],max_birth[d],num=dim[d])):
            #print(th, np.sum(np.logical_and(pds[:,0] < th, pds[:,1] > th)))
            res[-1][i] = np.sum(mlife[np.logical_and(pds[:,0] < th, pds[:,1] > th)])

    return np.sqrt(np.concatenate(res))

def comp_persistence_histogram(ph, num_bins, min_birth=None, max_birth=None, max_life=None, bandwidth=1):
#    print(args.num_bins)
    pds =[ph[ph[:,0] == i, 1:3] for i in range(2)]
    #print(len(pds[0]),len(pds[1]))
    life = [pd[:,1]-pd[:,0] for pd in pds]
    life = [l[l<1e+10] for l in life]  # remove permanent cycle
    life = [np.clip(l,0,max_life[i]) for i,l in enumerate(life)]
    # hsb = np.zeros(args.num_bins[2])
    # for k,ind in enumerate(np.searchsorted(lsb,pds[1][:,0])):
    #     hsb[ind] += (pds[1][k,1]-pds[1][k,0]) ## lifetime weighted count
    # hs0 = gaussian_kde(life[0])(ls) * len(life[0])
    # hs1 = gaussian_kde(life[1])(ls) * len(life[1])
    ## histogram for lifetime for each dimension
    hsl0, _ = np.histogram(life[0],bins=num_bins[0], range=(0,max_life[0]))
    hsl1, _ = np.histogram(life[1],bins=num_bins[1], range=(0,max_life[1]))
    # plt.hist(pds[0][:,0],weights=pds[0][:,1]-pds[0][:,0])
    # plt.show()
    # plt.hist(pds[1][:,0],weights=pds[1][:,1]-pds[1][:,0])
    # plt.show()
    ## histogram for birthtime for each dimension
    hsb0, _ = np.histogram(pds[0][:,0],bins=num_bins[2], range=(min_birth[0],max_birth[0]), weights=pds[0][:,1]-pds[0][:,0])
    hsb1, _ = np.histogram(pds[1][:,0],bins=num_bins[3], range=(min_birth[1],max_birth[1]), weights=pds[1][:,1]-pds[1][:,0])
    # lifetime weighting
    hsl0 = hsl0*(1.0+np.linspace(0,max_life[0],num_bins[0]))
    hsl1 = hsl1*(1.0+np.linspace(0,max_life[1],num_bins[1]))
    # smoothing
    hsl0 = kern_smooth(hsl0, bandwidth=bandwidth, kern='hanning')
    hsl1 = kern_smooth(hsl1, bandwidth=bandwidth, kern='hanning')
    hsb0 = kern_smooth(hsb0, bandwidth=bandwidth, kern='hanning')
    hsb1 = kern_smooth(hsb1, bandwidth=bandwidth, kern='hanning')
    # log and scaling
    hsl0,hsl1,hsb0,hsb1 = np.log(hsl0+1), np.log(hsl1+1), np.log(hsb0+1)/100, np.log(hsb1+1)
    # print(np.min(pds[0][:,0]),np.max(pds[0][:,0]),np.min(pds[1][:,0]),np.max(pds[1][:,0]))
    # print(np.min(pds[0][:,1]),np.max(pds[0][:,1]),np.min(pds[1][:,1]),np.max(pds[1][:,1]),"\n")
    return(np.concatenate([hsl0,hsl1,hsb0,hsb1]))

# kern = ['flat', 'hanning', 'hamming', 'bartlett', 'blackman']
def kern_smooth(y, bandwidth=11, kern='flat'):
    if bandwidth<2:
        return(y)
    b = int(bandwidth)
    if kern == 'flat':
        w=np.ones(bandwidth,'d')
    else:
        w=getattr(np,kern)(b)
    res = np.convolve(w/w.sum(),np.r_[y[b-1:0:-1],y,y[-2:-b-1:-1]],mode='valid')
    c = (len(res)-len(y))//2
    return(res[c:(c+len(y))])

if __name__== "__main__":

    args = arguments(mode="PHdict")

    grad = "grad" if args.gradient else ""
    if args.output is None:
        dn1,dn2 = os.path.split((os.path.normpath(args.target_dir))) # the leaf name
        phdn = os.path.join(dn1,"PH{}_{}_{}".format(grad,args.filtration,dn2))
        # if os.path.isdir(phdn):
        #     print("Please specify output directory!")
        #     exit()
        # else:
        args.output = phdn
        print("output will be saved under: ", args.output)

    ###
    print(args)
    target_dir = args.target_dir
    os.makedirs(args.output, exist_ok=True)

    with open(os.path.join(args.output, "args.json"), mode="w") as f:
        json.dump(args.__dict__, f, indent=4)
    gfns = []
    imgtypes = ['png','PNG','jpg','JPG','tif','TIF','tiff','TIFF']
    for it in imgtypes:
        gfns.extend(glob.glob(os.path.join(target_dir,"**/*.{}".format(it)), recursive=True))
    fns=sorted(list(set(gfns)))

    if args.label_type_pt == "persistence_histogram":
        print("compute and save persistence histogram...")
        meanPHl0 = np.zeros(args.num_bins[0])
        meanPHl1 = np.zeros(args.num_bins[1])
        meanPHb0 = np.zeros(args.num_bins[2])
        meanPHb1 = np.zeros(args.num_bins[3])
        for fname in tqdm(fns, total=len(fns)):
            bfn = os.path.splitext(os.path.basename(fname))[0]
            if args.imgtype=="npy":
                ph = np.load(fname)
            else:
                sample = np.array(Image.open(fname).convert('L'),dtype=np.float64)
                ph = comp_PH(sample, gradient=args.gradient, filtration=args.filtration)
                np.save(os.path.join(args.output, bfn), ph)
            hs = comp_persistence_histogram(ph, args.num_bins, min_birth=args.min_birth, max_birth=args.max_birth, max_life=args.max_life, bandwidth=args.bandwidth)
            np.save(os.path.join(args.output, bfn+"_hist"), hs.astype(np.float32))
            c1 = hs[:args.num_bins[0]]
            c2 = hs[args.num_bins[0]:(args.num_bins[0]+args.num_bins[1])]
            c3 = hs[(args.num_bins[0]+args.num_bins[1]):(args.num_bins[0]+args.num_bins[1]+args.num_bins[2])]
            c4 = hs[(args.num_bins[0]+args.num_bins[1]+args.num_bins[2]):]
            meanPHl0 += c1
            meanPHl1 += c2
            meanPHb0 += c3
            meanPHb1 += c4
            if args.save_fig:
                sns.lineplot(x=np.arange(len(c1)),y=c1, legend="full")
                sns.lineplot(x=np.arange(len(c2)),y=c2, legend="full",style=True, dashes=[(2,2)])
                sns.lineplot(x=np.arange(len(c3)),y=c3, legend="full",linewidth=2.5)
                sns.lineplot(x=np.arange(len(c4)),y=c4, legend="full",style=True, dashes=[(2,2)],linewidth=2.5)
                plt.savefig(os.path.join(args.output, bfn+"_histCurve.jpg"))
                plt.close()

        meanPHl0 /= len(fns)
        meanPHl1 /= len(fns)
        meanPHb0 /= len(fns)
        meanPHb1 /= len(fns)
        print(meanPHl0.max(), meanPHl1.max(),meanPHb0.max(),meanPHb1.max())
        print(sum(meanPHl0>0), sum(meanPHl1>0), sum(meanPHb0>0), sum(meanPHb1>0))
        if args.save_fig:
            sns.lineplot(x=np.arange(len(meanPHl0)),y=meanPHl0, legend="full")
            sns.lineplot(x=np.arange(len(meanPHl1)),y=meanPHl1, legend="full",style=True, dashes=[(2,2)])
            sns.lineplot(x=np.arange(len(meanPHb0)),y=meanPHb0, legend="full",linewidth=2.5)
            sns.lineplot(x=np.arange(len(meanPHb1)),y=meanPHb1, legend="full",style=True, dashes=[(2,2)],linewidth=2.5)
            plt.show()
    elif args.label_type_pt == "persistence_betticurve":
        print("compute and save betti curve...")
        meanPHl0 = np.zeros(args.num_bins[0])
        meanPHl1 = np.zeros(args.num_bins[1])
        for fname in tqdm(fns, total=len(fns)):
            sample = np.array(Image.open(fname).convert('L'),dtype=np.float64)
            ph = comp_PH(sample, gradient=args.gradient, filtration=args.filtration)
            res = comp_betticurve(ph, args.num_bins, min_birth=args.min_birth, max_birth=args.max_birth, max_life=args.max_life)
            bfn = os.path.splitext(os.path.basename(fname))[0]
            np.save(os.path.join(args.output, bfn+"_bettiCurve"), res.astype(np.float32))
            c1 = res[:args.num_bins[0]]
            c2 = res[args.num_bins[0]:]
            if args.save_fig:
                sns.lineplot(x=np.arange(len(c1)),y=c1, legend="full")
                sns.lineplot(x=np.arange(len(c2)),y=c2, legend="full",style=True, dashes=[(2,2)])
                plt.savefig(os.path.join(args.output, bfn+"_bettiCurve.jpg"))
                plt.close()
            meanPHl0 += c1
            meanPHl1 += c2
        meanPHl0 /= len(fns)
        meanPHl1 /= len(fns)
        print(meanPHl0.max(), meanPHl1.max())
        if args.save_fig:
            sns.lineplot(x=np.arange(len(meanPHl0)),y=meanPHl0, legend="full")
            sns.lineplot(x=np.arange(len(meanPHl1)),y=meanPHl1, legend="full",style=True, dashes=[(2,2)])
            plt.show()
    elif args.label_type_pt == "persistence_image":
        print("compute and save persistence images...")
        for fname in tqdm(fns, total=len(fns)):
            bfn = os.path.splitext(os.path.basename(fname))[0]
            img = Image.open(fname).convert('L')
            img = np.array(img, dtype=np.float64)
            ph = comp_PH(img, gradient=args.gradient, img_size=args.img_size, filtration=args.filtration)
            pims = comp_persistence_image(ph, args)
            np.save(os.path.join(args.output, bfn+"_persImg"), np.concatenate(pims).astype(np.float32))
            if args.save_fig:
                sns.lineplot(x=np.arange(len(pims[0])),y=pims[0], legend="full")
                sns.lineplot(x=np.arange(len(pims[1])),y=pims[1], legend="full",style=True, dashes=[(2,2)])
                plt.savefig(os.path.join(args.output, bfn+"_persImg.jpg"))
                plt.close()
                # plt.imshow(pims[0].reshape(10,5))
                # plt.savefig(os.path.join(args.output, bfn+"_persImg0.jpg"))
                # plt.close()
                # plt.imshow(pims[1].reshape(10,5))
                # plt.savefig(os.path.join(args.output, bfn+"_persImg1.jpg"))
                # plt.close()
    elif args.label_type_pt == "grid":
        from mpl_toolkits.axes_grid1 import make_axes_locatable
        fig = plt.figure(figsize=(21,10),tight_layout=True)
        n = min(len(fns),6)
        axes = fig.subplots(n, 6)
        for i in tqdm(range(n)):
            fname = fns[i]
            sample = np.array(Image.open(fname))
            mask = preprocess_image(sample, gradient=args.gradient, filtration='binarise', img_size=args.img_size) ## used only for preview
            dt = preprocess_image(sample, gradient=args.gradient, filtration=args.filtration, img_size=args.img_size)
            print(dt.min(),dt.max())
            ph = comp_PH(sample, gradient=args.gradient, filtration=args.filtration, img_size=args.img_size)
            axes[i,0].imshow(sample)
            axes[i,0].set_axis_off()
            axes[i,0].set_title(os.path.basename(fns[i]))
            axes[i,1].imshow(mask)
            axes[i,1].set_axis_off()
            im2 = axes[i,2].imshow(dt,vmin=args.min_birth[0],vmax=args.max_birth[0], )
            axes[i,2].set_axis_off()
            divider = make_axes_locatable(axes[i,2])
            cax = divider.append_axes('right', size='5%', pad=0.05)
            fig.colorbar(im2, cax=cax, orientation='vertical')

            axes[i,3].set_title("LC")
            res = comp_betticurve(ph, args.num_bins, min_birth=args.min_birth, max_birth=args.max_birth, max_life=args.max_life)
            #print(ph)
            sns.lineplot(x=np.arange(args.num_bins[0]),y=res[:args.num_bins[0]], ax=axes[i,3])
            sns.lineplot(x=np.arange(args.num_bins[1]),y=res[args.num_bins[0]:], style=True, dashes=[(2,3)], ax=axes[i,3])
            axes[i,4].set_title("HS")
            res = comp_persistence_histogram(ph, args.num_bins, min_birth=args.min_birth, max_birth=args.max_birth, max_life=args.max_life)
            sns.lineplot(x=np.arange(args.num_bins[0]),y=res[:args.num_bins[0]], ax=axes[i,4])
            sns.lineplot(x=np.arange(args.num_bins[1]),y=res[args.num_bins[0]:(args.num_bins[0]+args.num_bins[1])], style=True, dashes=[(2,2)], ax=axes[i,4])
            sns.lineplot(x=np.arange(args.num_bins[2]),y=res[(args.num_bins[0]+args.num_bins[1]):(args.num_bins[0]+args.num_bins[1]+args.num_bins[2])],linewidth=2.5, ax=axes[i,4])
            sns.lineplot(x=np.arange(args.num_bins[3]),y=res[(args.num_bins[0]+args.num_bins[1]+args.num_bins[2]):], style=True, dashes=[(2,2)],linewidth=2.5, ax=axes[i,4])
            axes[i,5].set_title("PI")
            res = comp_persistence_image(ph, args)
            #print(res[0].shape)
            sns.lineplot(x=np.arange(len(res[0])),y=res[0], ax=axes[i,5])
            sns.lineplot(x=np.arange(len(res[0])),y=res[1], style=True, dashes=[(2,2)], ax=axes[i,5])
            for ax in axes[i]:
                ax.legend([],[], frameon=False)
        plt.savefig(os.path.join(args.output,"persistence_vectors.jpg"))
        plt.show()

    ## compute persistence diagrams only
    else:
        print("compute and save persistent homology...")
        task = partial(comp_save_PH, args=args)
        with Pool(args.num_workers) as pool:
            with tqdm(total=len(fns), ascii=True, ncols=100) as t:
                for _ in pool.imap_unordered(task, fns):
                    t.update(1)


import cv2 as cv
import numpy as np
# method for optical flow extraction between two consequtive frames
def extract_the_optical_flow(frame_1, frame_2):
    # convert the images to gray scale
    prvs = cv.cvtColor(frame_1, cv.COLOR_BGR2GRAY)
    next = cv.cvtColor(frame_2, cv.COLOR_BGR2GRAY)

    # extracting the dense optical flow
    flow = cv.calcOpticalFlowFarneback(prvs, next, None, 0.5, 3, 15, 3, 5, 1.2, 0)

    # splitting to the x and y axis optical flow images
    gray_x = np.array(flow[..., 0])
    gray_x_to_write = cv.convertScaleAbs(gray_x, alpha=255.0)
    gray_y = np.array(flow[..., 1])
    gray_y_to_write = cv.convertScaleAbs(gray_y, alpha=255.0)
    return gray_x_to_write, gray_y_to_write


# David Ziama
# Class ID 3
# Write a Python program to create a symmetric difference.
# July 6, 2017

# The symmetric difference of two sets is the set of elements that are unique to each set

# Get items list from user.
# In this case I decided to use a person an their friends three favorite countries

user1 = input("What is User 1 name: ")
user2 = input("what is User 2 name: ")
print("\n")
print("Welcome to my symmetric difference finder", user1, "and", user2, ".")
print("This program will ask each of you for your favorite countries and see if you have any in commom")
print("\n")
one = input("User 1, What is your first favorite country? ")
two = input("User 1, What is your second favorite country? ")
three = input("User 1, What is your third favorite country? ")
four = input("User 2, What is your first favorite country? ")
five = input("User 2, What is your second favorite country? ")
six = input("User 2, What is your third favorite country? ")

# Define the sets

user1_list = set([one, two, three])
user2_list = set([four, five, six])

# Set symmetric difference
sym_diff = user1_list ^ user2_list

# Print the symmetric difference
print("The Symmetric Diffrence of" , user1, "and", user2, "favorite countries is" ,sym_diff)


# Using Recursion
def recur_factorial(n):
    if n == 1:
        return n
    else:
        return n*recur_factorial(n-1)
print(recur_factorial(5))

#Normal way
num = int(input("Enter a number: ") )
factorial = 1
if num < 0:
    print("Sorry, factorial does not exist for negative numbers")
elif num == 0:
    print("The factorial of 0 is 1")
else:
    for i in range (1,int(num)+1):
        factorial = factorial * i
print("Factorail of ",num , " is : ",factorial)


import csv
def delete_note():
    key = input("Введите ID заметки что бы выбрать её: ")
    new_data = []
    count = 0
    file = open("data.csv", 'r')
    data_text = csv.reader(file)
    for item in data_text:
        new_data.append(item)
        if ("id: " + key) in item:
            count +=1
            print(item)
    if count == 0:
        print("Ничего не найдено\n")
        return
    file.close()
    com = input("Удалить заметку? - y/n \n")
    if com == 'y' or com == 'Y':
        f = open("data.csv", 'w')
        writer = csv.writer(f)
        for item in new_data:
            if ("id: " + key) not in item:
                writer.writerow(item)
        print("Заметка удалена")
        f.close()
    elif com == 'n' or com == 'N':
        return
    else:
        print("Не корректная команда")


"""
Cosmic Ray General
==================
Code to model the evolution of the energy and position of cosmic ray protons emitted into the early universe IGM.
Essentially a set of common utility functions needed to compute the cosmic ray C coefficients and window functions.
"""

# Required modules
from typing import Union
import numpy as np
import scipy.constants
import os
from scipy.io import loadmat
from numba import njit
from scipy.special import gamma, hyp2f1

# Cosmology constants
H0 = 100*0.6704/3.086e19  # s^{-1}
OMEGA_M0 = 0.31687
OMEGA_B0 = 0.04902

fHe = 0.0732
H_FRAC = 1 - 4*fHe/(4*fHe + (1-fHe))

# Atomic Physics Constants
PROTON_M = 938.28  # In MeV/c^2
C_IN_MPC = scipy.constants.c / 3.086e22  # In Mpc s^{-1}
HI_ION_TIME_SCALE = 3e3 * 3.15e7  # From Sazonov & Sunyaev 2015, in s

# Derived Parameters
RHO_CRIT = 3 * H0**2 / (8 * np.pi * scipy.constants.G * 1e6)  # In kg cm^{-3}
TAU_INTERM = HI_ION_TIME_SCALE * scipy.constants.m_p / (RHO_CRIT * H_FRAC * OMEGA_B0)  # In s
KE_DECAY_CONST = 1**(3/2)/(H0*np.sqrt(OMEGA_M0)*TAU_INTERM)  # In MeV^{3/2}

# Model Parameters
KE_REL_TRANS_MEV = 100  # Kinetic energy of the relativistic to non-relativistic transition in MeV

# Simulation Parameters
LPIX = 3  # comoving Mpc

# General Utility Functions
@njit(cache=True)
def power_law_spectrum(ke: float, alpha: float = -2, ke_min: float = 1e-3, ke_max: float = 1e8) -> float:
    """ Probability density of a truncated power-law distribution normalized to a total of 1 MeV of energy

    Parameters
    ----------
    ke : Float
        Kinetic energy to evaluate at in MeV
    alpha : Float
        Exponent of the power-law, defaults to -2
    ke_min : Float
        Minimum KE cutoff of the power-law in MeV, defaults to 1e-3
    ke_max : Float
        Maximum KE cutoff of the power-law in MeV, defaults to 1e8

    Returns
    -------
    Float
        Normalized probability density in MeV^-2
    """
    if alpha == -2:
        nf = 1/(np.log(ke_max) - np.log(ke_min))
    else:
        nf = (alpha + 2)/(ke_max**(alpha + 2) - ke_min**(alpha + 2))
    return nf * (ke**alpha) * (ke <= ke_max) * (ke >= ke_min)

# Heating Functions
@njit(cache=True)
def normalized_heating_rate(ke: float, zs: float) -> float:
    """ Compute heating rate (normalized to remove fheat factor) of the IGM by CR protons by excitation and ionization

    Parameters
    ----------
    ke : Float
        Kinetic energies in MeV of the proton
    zs : Float
        Redshift heating rate is being evaluated at

    Returns
    -------
    Float
        Heating rate in MeV per redshift of the given proton
    """
    return KE_DECAY_CONST * ((1 + zs) ** (1/2)) * (ke ** (-1/2))

# CR Energy Evolution Functions
@njit(cache=True)
def initial_ke_rel(ke_now: float, z_now: float, z_ini: float) -> float:
    """ Initial kinetic energy of a CR proton if always in the relativistic regime

    Parameters
    ----------
    ke_now : Float
        Kinetic energy of protons now in MeV
    z_now : Float
        Redshift now
    z_ini : Float
        Redshift when proton was emitted

    Returns
    -------
    Float
        Kinetic energy of the proton when emitted in MeV
    """
    p_sq_factor = (ke_now**2 + 2*ke_now*PROTON_M) * (((1 + z_ini)/(1 + z_now))**2)
    return -PROTON_M + np.sqrt(PROTON_M**2 + p_sq_factor)

@njit(cache=True)
def current_ke_rel(ke_ini: float, z_now: float, z_ini: float) -> float:
    """ Current kinetic energy of a CR proton if always in the relativistic regime

    Parameters
    ----------
    ke_ini : Float
        Initial kinetic energy of the proton in MeV
    z_now : Float
        Redshift now
    z_ini : Float
        Redshift when proton was emitted

    Returns
    -------
    Float
        Kinetic energy of the proton at z_now in MeV
    """
    p_sq_factor = (ke_ini**2 + 2*ke_ini*PROTON_M) * (((1 + z_now)/(1 + z_ini))**2)
    return -PROTON_M + np.sqrt(PROTON_M**2 + p_sq_factor)

@njit(cache=True)
def initial_ke_non_rel(ke_now: float, z_now: float, z_ini: float) -> float:
    """ Initial kinetic energy of a CR proton if always in the non-relativistic regime

    Parameters
    ----------
    ke_now : Float
        Kinetic energy of protons now in MeV
    z_now : Float
        Redshift now
    z_ini : Float
        Redshift when proton was emitted

    Returns
    -------
    Float
        Kinetic energy of the proton when emitted in MeV
    """
    adb_factor = (1 + z_ini)**2 / (1 + z_now)**2
    return adb_factor * (ke_now**(3/2) + KE_DECAY_CONST*((1+z_now)**(3/2) -
                                                         ((1 + z_now)**3)*((1 + z_ini)**(-3/2))))**(2/3)

@njit(cache=True)
def current_ke_non_rel(ke_ini: float, z_now: float, z_ini: float) -> float:
    """ Current kinetic energy of a CR proton if always in the non-relativistic regime

    Parameters
    ----------
    ke_ini : Float
        Initial kinetic energy of the proton in MeV
    z_now : Float
        Redshift now
    z_ini : Float
        Redshift when proton was emitted

    Returns
    -------
    Float
        Kinetic energy of the proton at z_now in MeV
    """
    adb_factor = (1 + z_now)**2 / (1 + z_ini)**2
    return adb_factor * (ke_ini**(3/2) - KE_DECAY_CONST*(((1+z_now)**(-3/2))*((1 + z_ini)**3) -
                                                         (1 + z_ini)**(3/2)))**(2/3)

@njit(cache=True)
def transition_redshifts(ke_now: float, z_now: float) -> float:
    """ Compute redshift a proton would have transitioned between the relativistic and non-relativistic regime

    Parameters
    ----------
    ke_now : Float
        Kinetic energy of proton now in MeV
    z_now : Float
        Redshift now

    Returns
    -------
    Float
        Redshift of relativistic to non-relativistic transition
    """
    ke_parameter = KE_DECAY_CONST*((1 + z_now)**(-3/2)) + (ke_now**(3/2))/((1 + z_now)**3)
    return ((-KE_DECAY_CONST + np.sqrt(KE_DECAY_CONST**2 + 4*ke_parameter*(KE_REL_TRANS_MEV**(3/2)))) /
            (2*(KE_REL_TRANS_MEV**(3/2))))**(-2/3) - 1

@njit(cache=True)
def initial_ke(ke_now: float, z_now: float, z_ini: float) -> float:
    """ Initial kinetic energy of an arbitrary CR proton

    Parameters
    ----------
    ke_now : Float
        Kinetic energy of proton now in MeV
    z_now : Float
        Redshift now
    z_ini : Float
        Redshift when proton was emitted

    Returns
    -------
    Float
        Kinetic energy of the proton when emitted in MeV
    """

    # Assume initially no transitions occur
    if ke_now > KE_REL_TRANS_MEV:
        ini_ke = initial_ke_rel(ke_now, z_now, z_ini)
    else:
        ini_ke = initial_ke_non_rel(ke_now, z_now, z_ini)
        if ini_ke > KE_REL_TRANS_MEV:  # Check if transition occurs
            ini_ke = initial_ke_rel(KE_REL_TRANS_MEV, transition_redshifts(ke_now, z_now), z_ini)
    return ini_ke

@njit(cache=True)
def current_ke(ke_ini: float, z_now: float, z_ini: float) -> float:
    """ Current kinetic energy of an arbitrary CR proton

    Parameters
    ----------
    ke_ini : Float
        Initial kinetic energy of the proton in MeV
    z_now : Float
        Redshift now
    z_ini : Float
        Redshift when proton was emitted

    Returns
    -------
    Float
        Kinetic energy of the proton at z_now in MeV
    """

    # If at a non-relativistic initial energy then it stays non-relativistic
    if ke_ini <= KE_REL_TRANS_MEV:
        ke_now = current_ke_non_rel(ke_ini, z_now, z_ini)

    # Otherwise, starts relativistic and might transition to non-relativistic later
    else:
        ke_now = current_ke_rel(ke_ini, z_now, z_ini)
        if ke_now < KE_REL_TRANS_MEV:  # Check if transition occurs
            z_trans = (1 + z_ini) * (((KE_REL_TRANS_MEV**2 + 2*KE_REL_TRANS_MEV*PROTON_M) /
                                      (ke_ini**2 + 2*ke_ini*PROTON_M))**(1/2)) - 1
            ke_now = current_ke_non_rel(KE_REL_TRANS_MEV, z_now, z_trans)

    return ke_now

@njit(cache=True)
def jacobian_factor(ke_now: float, z_now: float, z_ini: float) -> float:
    """ Jacobian factor that transforms between density of KE now and KE at emission

    Branch-less implementation for numpy arrays

    Parameters
    ----------
    ke_now : Float
        Kinetic energy of proton now in MeV
    z_now : Float
        Redshift now
    z_ini : Float
        Redshift when proton was emitted

    Returns
    -------
    Float
        Jacobian factor (dimensionless) dT'/dT
    """

    # Get initial KE of proton
    ini_ke = initial_ke(ke_now, z_now, z_ini)

    # Assume initially no transitions occur
    if ke_now > KE_REL_TRANS_MEV:
        jacob = ((ke_now + PROTON_M)/(ini_ke + PROTON_M))*(((1 + z_ini) / (1 + z_now))**2)
    else:
        if ini_ke < KE_REL_TRANS_MEV:
            jacob = np.sqrt(ke_now / ini_ke)*(((1 + z_ini) / (1 + z_now))**3)
        else:  # Transitioning case
            jacob = ((KE_REL_TRANS_MEV + PROTON_M) / (ini_ke + PROTON_M)) * np.sqrt(ke_now / KE_REL_TRANS_MEV) * \
                    (((1 + z_ini) ** 2) * (1 + transition_redshifts(ke_now, z_now)) / ((1 + z_now) ** 3))
    return jacob

# CR Path Length Functions
def g_1_func(x: np.ndarray) -> np.ndarray:
    """ Special function required to compute path length in the non-relativistic regime

    G_1 == 2F1(-1/3, -1/3; 2/3; x)

    Parameters
    ----------
    x : Array
        Arguments to evaluate G_1 at

    Returns
    -------
    Array
        G_1(x)
    """

    # Input sanitization
    if np.any(x < 0):
        raise ValueError(f'One or more of the given arguments is outside the range 0 <= x <= 1 for which '
                         f'this function is defined.')

    # Initialize data structure for the output
    output = np.empty(np.shape(x), dtype=float)

    # Consider the two special boundary cases
    output[x == 0] = 1
    output[x >= 1] = gamma(2/3)*gamma(4/3)

    # General case
    output[np.logical_and(x > 0, x < 1)] = hyp2f1(-1/3, -1/3, 2/3, x[np.logical_and(x > 0, x < 1)])
    return output

def g_2_func(x: np.ndarray) -> np.ndarray:
    """ Special function required to compute path length in the relativistic regime

        G_2 == 2F1(1/4, 1/2; 5/4; -x^2)      x in R

        Parameters
        ----------
        x : Array
            Arguments to evaluate G_2 at, x must be real

        Returns
        -------
        Array
            G_2(x)
        """
    return hyp2f1(1/4, 1/2, 5/4, -np.power(x, 2))

def alpha_1_parameter(ke_0: np.ndarray, z_0: float) -> np.ndarray:
    """ Dimensionless parameter required to compute path length in the non-relativistic regime

    alpha_1 == ((H_0 sqrt{Omega_m} tau_{inter} ke_0^{3/2})/((1 MeV)^{3/2}(1+z_0)^3)  + (1+z_0)^{-3/2})^{-2/3}

    Equivalently alpha_1 can be interpreted as 1 + z_{abs}, the inverse scale factor when a cosmic ray proton would
    be thermalized with the IGM.

    Parameters
    ----------
    ke_0 : Array
        Initial KE of the protons in MeV
    z_0: Float
        Emission redshift of the proton

    Returns
    -------
    Array
        alpha_1 parameter evaluated for each of these protons
    """
    return (((ke_0**(3/2)) * ((1 + z_0)**(-3)))/KE_DECAY_CONST + (1 + z_0)**(-3/2))**(-2/3)

def alpha_2_parameter(ke_0: np.ndarray, z_0: float) -> np.ndarray:
    """ Dimensionless parameter required to compute path length in the relativistic regime

    alpha_2 == sqrt(T_0^2/(m_p^2 c^4) + 2 T_0 / (m_pc^2)) / (1 + z_0)

    Parameters
    ----------
    ke_0 : Array
        Initial KE of the protons in MeV
    z_0: Float
        Emission redshift of the proton

    Returns
    -------
    Array
        alpha_2 parameter evaluated for each of these protons
    """
    norm_ke = ke_0 / PROTON_M
    return np.sqrt(norm_ke**2 + 2*norm_ke) / (1 + z_0)

def comoving_path_length_non_rel(z: float, ke_0: Union[float, np.ndarray], z_0: Union[float, np.ndarray]) -> np.ndarray:
    """ Comoving path length a non-relativistic proton would travel between z_0 and z

    Note for performance it is assumed that none of the protons have been absorbed, any such protons should be filtered
    out before calling this function

    Parameters
    ----------
    z : Float
        Redshift proton is at now
    ke_0 : Array
        Initial kinetic energies of the protons in MeV
    z_0 : Float
        Source redshift of protons

    Returns
    -------
    Array
        Comoving path lengths travelled in comoving Megaparsec
    """
    # Pre-computations
    alpha_1 = alpha_1_parameter(ke_0, z_0)
    prefactor = (C_IN_MPC/H0) * ((8/(OMEGA_M0 * PROTON_M))**(1/2)) * (KE_DECAY_CONST**(1/3)) * (alpha_1**(-1/2))

    # Primary computation
    output = ((1 + z_0)**(1/2))*g_1_func((alpha_1/(1 + z_0))**(3/2)) - \
             ((1 + z)**(1/2))*g_1_func((alpha_1/(1 + z))**(3/2))
    return prefactor*output

def comoving_path_length_rel(z: float, ke_0: np.ndarray, z_0: float) -> np.ndarray:
    """ Comoving path length a relativistic proton would travel between z_0 and z

    Parameters
    ----------
    z : Float
        Redshift proton is at now
    ke_0 : Array
        Initial kinetic energies of the protons in MeV
    z_0 : Float
        Source redshift of protons

    Returns
    -------
    Array
        Comoving path lengths travelled in comoving Megaparsec
    """
    # Pre-computations
    alpha_2 = alpha_2_parameter(ke_0, z_0)
    prefactor = 2 * C_IN_MPC * alpha_2 / (H0 * np.sqrt(OMEGA_M0))

    # Primary computation
    output = np.sqrt(1 + z_0)*g_2_func(alpha_2*(1 + z_0)) - np.sqrt(1 + z)*g_2_func(alpha_2*(1 + z))
    return prefactor*output

def comoving_path_proton(z: float, ke_0: np.ndarray, z_0: float, rel_only: bool = False) -> np.ndarray:
    """ Comoving path length a CR proton would travel between z_0 and z

    Note for performance it is assumed that none of the protons have been absorbed, any such protons should be filtered
    out before calling this function

    Parameters
    ----------
    z : Float
        Redshift proton is at now
    ke_0 : Array
        Initial kinetic energies of the protons in MeV
    z_0 : Float
        Source redshift of protons
    rel_only : Boolean
        Flag to only use the relativistic regime solution rather than the full solution

    Returns
    -------
    Array
        Comoving path lengths travelled in comoving Megaparsec
    """
    # Datastructure for output
    output = np.empty(np.shape(ke_0), dtype=float)

    # Handle relativistic only case
    if rel_only:
        output = comoving_path_length_rel(z, ke_0, z_0)
        return output

    # Otherwise, the simplest case is if starts non-relativistic, then as they only lose energy it stays
    # non-relativistic
    output[ke_0 <= KE_REL_TRANS_MEV] = comoving_path_length_non_rel(z, ke_0[ke_0 <= KE_REL_TRANS_MEV], z_0)

    # For relativistic case we have the added complication that the proton can change into the non-relativistic regime
    # if sufficient time has passed, this transition would occur at a redshift
    z_rel_trans = (1 + z_0) * np.sqrt((KE_REL_TRANS_MEV ** 2 + 2 * KE_REL_TRANS_MEV * PROTON_M)
                                      / (ke_0 ** 2 + 2 * ke_0 * PROTON_M)) - 1

    # If z >= z_rel_trans then this transition has not occurred yet, and so we do not have to deal with it
    output[np.logical_and(ke_0 > KE_REL_TRANS_MEV, z >= z_rel_trans)] = \
        comoving_path_length_rel(z, ke_0[np.logical_and(ke_0 > KE_REL_TRANS_MEV, z >= z_rel_trans)], z_0)

    # If z < z_rel_trans then we do need to handle the transition by summing the seperated Delta Rs for each part of the
    # protons journey
    slice_trans_p = np.logical_and(ke_0 > KE_REL_TRANS_MEV, z < z_rel_trans)
    output[slice_trans_p] = comoving_path_length_rel(z_rel_trans[slice_trans_p], ke_0[slice_trans_p], z_0) + \
        comoving_path_length_non_rel(z, KE_REL_TRANS_MEV, z_rel_trans[slice_trans_p])
    return output


"""
Make best efforts to decide if a task needs to run again, or if it is pointless
because the source code hasn't changed.

Sometimes task depends on the world, not the source code

Strategies
- folder or source tree changed
- file changed
- list of files changed

Is it worse is lint counting. Move to different module?
"""
import functools
import hashlib
import os
import shutil
import sys
import time
from typing import Any, Callable, Optional, TypeVar, cast

from checksumdir import dirhash

from navio_tasks import settings as settings
from navio_tasks.settings import PROJECT_NAME, SRC
from navio_tasks.utils import inform

# pylint: disable=invalid-name
FuncType = Callable[..., Any]
# pylint: disable=invalid-name
F = TypeVar("F", bound=FuncType)
CURRENT_HASH = None

# bash to find what has change recently
# find src/ -type f -print0 | xargs -0 stat -f "%m %N" | sort -rn | head -10 |
# cut -f2- -d" "

class BuildState:
    """
    Try not to re-do what doesn't need to be redone
    """

    def __init__(self, what: str, directory: str) -> None:
        """
        Set initial state
        """
        self.what = what
        self.directory = directory
        if not os.path.exists(f"{settings.CONFIG_FOLDER}/.build_state"):
            os.makedirs(f"{settings.CONFIG_FOLDER}/.build_state")
        self.state_file_name = (
            f"{settings.CONFIG_FOLDER}/.build_state/last_change_{what}.txt"
        )

    def oh_never_mind(self) -> None:
        """
        If a task fails, we don't care if it didn't change since last, re-run,
        """
        # noinspection PyBroadException
        # pylint: disable=bare-except
        try:
            os.remove(self.state_file_name)
        except:  # noqa: B001
            pass

    def has_source_code_tree_changed(self) -> bool:
        """
        If a task succeeds & is re-run and didn't change, we might not
        want to re-run it if it depends *only* on source code
        """
        # pylint: disable=global-statement
        global CURRENT_HASH
        directory = self.directory

        # if CURRENT_HASH is None:
        # inform("hashing " + directory)
        # inform(os.listdir(directory))
        CURRENT_HASH = dirhash(
            directory,
            "md5",
            ignore_hidden=True,
            # changing these exclusions can cause dirhas to skip EVERYTHING
            # excluded_files=[".coverage", "lint.txt"],
            excluded_extensions=[".pyc"],
        )

        inform("Searching " + self.state_file_name)
        if os.path.isfile(self.state_file_name):
            with open(self.state_file_name, "r+") as file:
                last_hash = file.read()
                if last_hash != CURRENT_HASH:
                    file.seek(0)
                    file.write(CURRENT_HASH)
                    file.truncate()
                    return True
                return False

        # no previous file, by definition not the same.
        with open(self.state_file_name, "w") as file:
            file.write(CURRENT_HASH)
            return True

def oh_never_mind(what: str) -> None:
    """
    If task fails, remove file that says it was recently run.
    Needs to be like this because tasks can change code (and change the hash)
    """
    state = BuildState(what, PROJECT_NAME)
    state.oh_never_mind()

def has_source_code_tree_changed(
    task_name: str, expect_file: Optional[str] = None
) -> bool:
    """
    Hash source code tree to know if it has changed

    Also check if an expected output file exists or not.
    """
    if expect_file:
        if os.path.isdir(expect_file) and not os.listdir(expect_file):
            os.path.dirname(expect_file)
            # output folder empty
            return True
        if not os.path.isfile(expect_file):
            # output file gone
            return True
    state = BuildState(task_name, os.path.join(SRC, PROJECT_NAME))
    return state.has_source_code_tree_changed()

def skip_if_no_change(name: str, expect_files: Optional[str] = None) -> F:
    """
    Don't run decorated task if nothing in the source has changed.
    """

    # https://stackoverflow.com/questions/5929107/decorators-with-parameters
    def real_decorator(func: F) -> F:
        """Wrapper"""

        @functools.wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Callable:
            """Wrapper"""
            if not has_source_code_tree_changed(name, expect_files):
                inform("Nothing changed, won't re-" + name)
                return lambda x: None
            try:
                return func(*args, **kwargs)
            except:  # noqa: B001
                oh_never_mind(name)
                raise

        return cast(F, wrapper)

    return cast(F, real_decorator)

def hash_it(path: str) -> str:
    """
    Hash a single file. Return constant if it doesn't exist.
    """
    if not os.path.exists(path):
        return "DOESNOTEXIST"
    with open(path, "rb") as file_handle:
        return hashlib.sha256(file_handle.read()).hexdigest()

def skip_if_this_file_does_not_change(name: str, file: str) -> F:
    """
    Skip decorated task if this referenced file didn't change. Useful
    if a task depends on a single file and not (potentially) any file in the source tree
    """

    def real_decorator(func: F) -> F:
        """Wrapper"""

        @functools.wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Callable:
            """Wrapper"""
            state_file = (
                f"{settings.CONFIG_FOLDER}/.build_state/file_hash_" + name + ".txt"
            )
            previous_hash = "catdog"
            if os.path.exists(state_file):
                with open(state_file) as old_file:
                    previous_hash = old_file.read()

            new_hash = hash_it(file)
            if new_hash == previous_hash:
                inform("Nothing changed, won't re-" + name)
                return lambda x: f"Skipping {name}, no change"
            if not os.path.exists(f"{settings.CONFIG_FOLDER}/.build_state"):
                os.makedirs(f"{settings.CONFIG_FOLDER}/.build_state")
            with open(state_file, "w+") as state:
                state.write(new_hash)
            try:

                return func(*args, **kwargs)
            except:  # noqa: B001
                # reset if step fails
                os.remove(state_file)
                raise

        return cast(F, wrapper)

    return cast(F, real_decorator)

def reset_build_state() -> None:
    """
    Delete all .build_state & to force all steps to re-run next build
    """
    if os.path.exists(f"{settings.CONFIG_FOLDER}/.build_state"):
        shutil.rmtree(f"{settings.CONFIG_FOLDER}/.build_state")
    if not os.path.exists(f"{settings.CONFIG_FOLDER}/.build_state"):
        os.makedirs(f"{settings.CONFIG_FOLDER}/.build_state")

def timed() -> F:
    """This decorator prints the execution time for the decorated function."""

    def real_decorator(func: F) -> F:
        """Wrapper"""

        @functools.wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Callable:
            """Wrapper"""
            start = time.time()
            result = func(*args, **kwargs)
            end = time.time()
            inform("{} ran in {}s".format(func.__name__, round(end - start, 2)))
            return result

        return cast(F, wrapper)

    return cast(F, real_decorator)

def is_it_worse(task_name: str, current_rows: int, margin: int) -> bool:
    """
    Logic for dealing with a code base with very large amounts of lint.

    You will never fix it all, just don't make it worse.
    """
    if not os.path.exists(f"{settings.CONFIG_FOLDER}/.build_state"):
        os.makedirs(f"{settings.CONFIG_FOLDER}/.build_state")
    file_name = f"{settings.CONFIG_FOLDER}/.build_state/last_count_{task_name}.txt"

    last_rows = sys.maxsize
    if os.path.isfile(file_name):
        with open(file_name, "r+") as file:
            last_rows = int(file.read())
            if last_rows != current_rows:
                file.seek(0)
                file.write(str(current_rows))
                file.truncate()

    return current_rows > (last_rows + margin)


from doubledqn import Agent
import gym
import numpy as np

env = gym.make('CartPole-v1')

agent = Agent(state_dims=4,n_actions=2)

scores = []
for i in range(5_000):
    done = False
    score = 0
    state = env.reset()
    while not done:
        action = agent.get_action(state)
        new_state, reward, done, _ = env.step(action)
        score += reward
        agent.store(state,action,reward,new_state,done)
        state = new_state
        agent.train(i)

    scores.append(score)
    avg_score = np.mean(scores[max(0, i - 100):(i+1)])
    print(f'episode {i} score %.2f {score} avg_score %.2f {avg_score} epsilon {agent.epsilon}')


# -*- coding: utf-8; py-indent-offset:4 -*-

import datetime as dt
from hiquant import *

def XXXtest_backtrade():
    # test back trade
    date_start = date_from_str('1 months ago')
    date_end = date_from_str('yesterday')
    market = Market(date_start, date_end)

    df = get_stockpool_df(['600036', '000002'])
    df.to_csv('output/mytest.csv', index= False)

    trader = Trader(market)
    fund = Fund(market, trader, 'fund_1', {
        'name': 'fund no.1',
        'start_cash': '1000000.00',
        'strategy': 'strategy/001_macd.py',
        'stock_pool': 'output/mytest.csv',
    })
    trader.add_fund(fund)
    market.set_verbose()
    trader.set_verbose()
    trader.run_fund(date_start, date_end)
    trader.print_report()
    trader.plot(compare_index= 'sh000300', out_file='output/test_trader.png')

def XXXtest_run():
    date_start = date_from_str('today')
    date_end = date_from_str('tomorrow')
    market = Market(date_start, date_end)
    market.set_force_open(True)

    trader = Trader(market)
    fund = Fund(market, trader, 'fund_1', {
        'name': 'fund no.1',
        'start_cash': '1000000.00',
        'strategy': 'strategy/001_macd.py',
        'stock_pool': 'output/mytest.csv',
    })
    trader.add_fund(fund)
    market.set_verbose()
    trader.set_verbose()
    trader.run_fund(date_start, date_start + dt.timedelta(seconds=10))
    #trader.print_report()
    #trader.plot(compare_index= 'sh000300', out_file='output/trader_run.png')