from cmath import isnan
from dataclasses import dataclass, field
from email.policy import default
import glob
from typing import Iterable, Type, Optional
import h5py
import logging
import json
import re
import math
import copy
import datetime
import os

import bisect
import kmz

from objects import *
from moos_messages import *
from pprint import pprint
from types import MethodType
from zipfile import *

# JAIA message types as python dataclasses
from jaia_messages import *

INT32_MAX = (2 << 30) - 1
UINT32_MAX = (2 << 31) - 1

LOG_DIR='logs'

# Utility functions

def set_directory(directory):
    global LOG_DIR
    LOG_DIR = directory
    logging.info(f'Log directory: {directory}')
    os.makedirs(directory, exist_ok=True)

def get_root_item_path(path, root_item=''):
    '''Get the path to a root_item associated with that path'''
    components = path.split('/')
    components = components[:2] + [root_item]
    return '/'.join(components)

def h5_get_series(dataset):
    '''Get a filtered, JSON-serializable representation of an h5 dataset'''
    dtype = dataset.dtype

    def from_float(x):
        x = float(x)
        if math.isnan(x):
            return None
        return x

    def from_int32(x):
        if x == INT32_MAX:
            return None
        return int(x)

    def from_uint32(x):
        if x == UINT32_MAX:
            return None
        return int(x)

    dtype_proc = {
        'f': from_float,
        'i': from_int32,
        'u': from_uint32
    }

    map_proc = dtype_proc[dtype.kind]
    filtered_list = [map_proc(x) for x in dataset]

    return filtered_list

def h5_get_string(dataset):
    s = ''.join([chr(a) for a in list(list(dataset)[0])])
    return s

def h5_get_hovertext(dataset):
    '''Get the hovertext for an h5 dataset'''

    # Get the enum value names
    try:
        enum_names = dataset.attrs['enum_names']
        enum_values = dataset.attrs['enum_values']
        enum_dict = { int(enum_values[index]): enum_names[index] for index in range(0, len(enum_values))}
        return enum_dict

    except KeyError:
        return None

def itemsmatching(file: h5py.File, regular_expression: re.Pattern):
    '''Returns an iterator for the paths matching regular_expression found in HDF5 file'''
    matching_items = []

    def func(name, object):
        m = regular_expression.match(name)

        if m is not None:
            matching_items.append(name)

    file.visititems(func)

    for item in matching_items:
        yield item

# Path descriptions

try:
    path_descriptions = json.load(open('jaiabot_paths.json'))
except FileNotFoundError:
    path_descriptions = {}

def jaia_get_description(path):
    for description in path_descriptions:
        if 'path' in description:
            if description['path'] == path:
                return description

        if 'path_regex' in description:
            if re.match(description['path_regex'], path):
                return description

    return None

def get_title_from_path(path):
    components = path.split('/')
    if len(components) < 2:
        logging.warning(f'Not enough components in path: {path}')
        return ''

    components = components[1:]

    message_type_components = components[0].split('.')

    if len(message_type_components) < 1:
        logging.warning(f'Invalid path: {path}')
        return ''

    components[0] = components[0].split('.')[-1]
    return '/'.join(components)

# Path regular expressions

BOT_STATUS_RE = re.compile(r'^jaiabot::bot_status.*;([0-9]+)/jaiabot.protobuf.BotStatus$')

# Data fetch functions

def get_logs():
    '''Get list of available logs'''
    results = []
    logging.warning('log_dir = ' + LOG_DIR)
    for filename in glob.glob(LOG_DIR + '/*_*_*.h5'):
        components = re.match(r'.*/(.+)_(.+)_(.+)\.h5$', filename)
        bot, fleet, date_string = components.groups()

        try:
            date = datetime.datetime.strptime(date_string, r'%Y%m%dT%H%M%S').replace(tzinfo=datetime.timezone.utc)
        except ValueError:
            logging.warning(f'No date in file {filename}')
            continue

        # Get duration of this log
        try:
            path = 'goby::health::report/goby.middleware.protobuf.VehicleHealth/_utime_'
            h5_file = h5py.File(filename)
            start = h5_file[path][0]
            end = h5_file[path][-1]
            duration = int(end - start)
        except (OSError, KeyError):
            logging.warning(f'No duration in file {filename}')
            continue

        results.append({
            'bot': bot,
            'fleet': fleet,
            'timestamp': date.timestamp(),
            'duration': duration,
            'filename': filename
        })

    return results

def delete_log(logName: str):
    logGlob = f'{LOG_DIR}/{logName}.*'
    logPaths = glob.glob(logGlob)

    for logPath in logPaths:
        logging.warning(f'Deleting file {logPath}')
        os.remove(logPath)

def get_fields(log_names, root_path='/'):
    '''Get a list of the fields below a root path in a set of logs'''
    fields = set()

    if log_names is None:
        return []

    if root_path is None or root_path == '':
        root_path = '/'
    else:
        # h5py doesn't like initial slashes, unless it's the root path
        root_path = root_path.lstrip('/')

    log_names = log_names.split(',')

    for log_name in log_names:
        h5_file = h5py.File(log_name)
        try:
            fields = fields.union(h5_file[root_path].keys())
        except AttributeError:
            # If this path is a dataset, it has no "keys()"
            pass

    series = list(fields)
    series.sort()

    return series

@dataclass
class Series:
    utime: list
    y_values: list
    hovertext: dict

    def __init__(self, log=None, path=None, scheme=1, invalid_values=set()) -> None:
        self.utime = []
        self.y_values = []
        self.hovertext = {}

        if log:
            try:
                _utime__array = log[get_root_item_path(path, '_utime_')]
                _scheme__array = log[get_root_item_path(path, '_scheme_')]
                path_array = log[path]

                series = zip(h5_get_series(_utime__array), h5_get_series(_scheme__array), h5_get_series(path_array))
                series = filter(lambda pt: pt[1] == scheme and pt[2] not in invalid_values, series)

                self.utime, schemes, self.y_values = zip(*series)
            except (ValueError, KeyError):
                logging.warning(f'No valid data found for log: {log.filename}, series path: {path}')
                self.utime = []
                self.schemes = []
                self.y_values = []
                self.hovertext = {}

                return

            self.hovertext = h5_get_hovertext(log[path]) or {}

    def __add__(self, other_series: 'Series'):
        r = copy.copy(self)
        r.utime += list(other_series.utime)
        r.y_values += list(other_series.y_values)
        r.hovertext.update(other_series.hovertext)
        return r

    def sort(self):
        if len(self.utime) > 0 and len(self.y_values) > 0:
            self.utime, self.y_values = zip(*sorted(zip(self.utime, self.y_values)))
        else:
            logging.warning(f'Not enough values to sort.  len(utime) = {len(self.utime)}, len(y_values) = {len(self.y_values)}')

    def getValueAtTime(self, t):
        index = bisect.bisect_left(self.utime, t)
        if index == 0:
            return None
        else:
            return self.y_values[index - 1]

def get_series(log_names, paths):
    '''Get a series'''

    if log_names is None or paths is None:
        return []

    series_list = []

    log_names = log_names.split(',')
    paths = [path.lstrip('/') for path in paths.split(',')]

    # Open all our logs
    logs = [h5py.File(log_name) for log_name in log_names]

    # Get the series from the logs
    for path in paths:
        series_description = jaia_get_description(path) or {}
        invalid_values = set(series_description.get('invalid_values', []))

        series = Series()

        for log in logs:
            try:
                series += Series(log=log, path=path, scheme=1, invalid_values=invalid_values)
            except KeyError as e:
                logging.warn(e)
                continue

        series.sort()

        title = get_title_from_path(path)
        y_axis_title = title
        units = series_description.get('units')

        if units is not None:
            y_axis_title += f'\n({units})'

        series_list.append({
            'path': path,
            'title': title,
            'y_axis_title': y_axis_title,
            '_utime_': series.utime,
            'series_y': series.y_values,
            'hovertext': series.hovertext
        })

    return series_list

def get_map(log_names):
    if log_names is None:
        return {}

    log_names = log_names.split(',')

    # Open all our logs
    logs = [h5py.File(log_name) for log_name in log_names]

    NodeStatus_lat_path = 'goby::middleware::frontseat::node_status/goby.middleware.frontseat.protobuf.NodeStatus/global_fix/lat'
    NodeStatus_lon_path = 'goby::middleware::frontseat::node_status/goby.middleware.frontseat.protobuf.NodeStatus/global_fix/lon'
    NodeStatus_heading_path = 'goby::middleware::frontseat::node_status/goby.middleware.frontseat.protobuf.NodeStatus/pose/heading'
    NodeStatus_course_over_ground_path = 'goby::middleware::frontseat::node_status/goby.middleware.frontseat.protobuf.NodeStatus/pose/course_over_ground'
    DesiredSetpoints_heading_path = 'jaiabot::desired_setpoints/jaiabot.protobuf.DesiredSetpoints/helm_course/heading'

    bot_id_to_map_series = {}

    for log in logs:

        bot_id_string: str = None

        # Find BotStatus to get bot_id
        for path in itemsmatching(log, BOT_STATUS_RE):
            bot_id_series = Series(log=log, path=path + '/bot_id', invalid_values=[None])
            bot_id_string = str(bot_id_series.y_values[0])
            break

        if bot_id_string is None:
            logging.warning(f'No bot_id for {log.filename}')
            continue

        map_series = bot_id_to_map_series.setdefault(bot_id_string, [])

        lat_series = Series(log=log, path=NodeStatus_lat_path, invalid_values=[0])
        lon_series = Series(log=log, path=NodeStatus_lon_path, invalid_values=[0])
        heading_series = Series(log=log, path=NodeStatus_heading_path)
        course_over_ground_series = Series(log=log, path=NodeStatus_course_over_ground_path)
        desired_heading_series = Series(log=log, path=DesiredSetpoints_heading_path)

        for i, lat in enumerate(lat_series.y_values):
            most_recent_desired_heading = desired_heading_series.getValueAtTime(lat_series.utime[i])
            map_series.append([
                lat_series.utime[i],
                lat_series.y_values[i],
                lon_series.y_values[i],
                heading_series.y_values[i],
                course_over_ground_series.y_values[i],
                most_recent_desired_heading
            ])

    return bot_id_to_map_series

HUB_COMMAND_RE = re.compile(r'jaiabot::hub_command.*;([0-9]+)')

def get_commands(log_filenames):

    # Open all our logs
    log_files = [h5py.File(log_name) for log_name in log_filenames]

    # A dictionary mapping bot_id to an array of mission dictionaries
    results = {}

    for log_file in log_files:

        # Search for Command items
        for path in log_file.keys():

            m = HUB_COMMAND_RE.match(path)
            if m is not None:
                hub_command_path = path

                command_path = hub_command_path + '/jaiabot.protobuf.Command'
                commands = jaialog_get_object_list(log_file[command_path], repeated_members={"goal"})

                if len(commands) > 0:
                    bot_id_string = str(commands[0]['bot_id'])

                    if bot_id_string not in results:
                        results[bot_id_string] = []

                    results[bot_id_string].extend(commands)

    # Sort each bot's command list by time of command issuance
    for bot_id_string in results:
        results[bot_id_string].sort(key=lambda command: command['_utime_'])

    return results

def get_active_goals(log_filenames):

    # Open all our logs
    log_files = [h5py.File(log_name) for log_name in log_filenames]

    # A dictionary mapping bot_id to an array of active_goal dictionaries
    results = {}

    for log_file in log_files:

        def visit_path(name, object):
            m = BOT_STATUS_RE.match(name)

            if m is not None:
                _utime_ = h5_get_series(log_file[name + '/_utime_'])
                bot_ids = h5_get_series(log_file[name + '/bot_id'])
                active_goals = h5_get_series(log_file[name + '/active_goal'])

                last_active_goal = None
                for i in range(len(_utime_)):
                    bot_id_string = str(bot_ids[i])

                    active_goal_list_for_this_bot = results.setdefault(bot_id_string, [])

                    try:
                        last_active_goal = active_goal_list_for_this_bot[-1]['active_goal']
                    except IndexError:
                        last_active_goal = None

                    if active_goals[i] != last_active_goal:
                        active_goal_list_for_this_bot.append({
                            '_utime_': _utime_[i],
                            'active_goal': active_goals[i]
                        })

        log_file.visititems(visit_path)

    return results

TASK_PACKET_RE = re.compile(r'jaiabot::task_packet.*;([0-9]+)')

def get_task_packet_dicts(log_filenames: List[str], scheme=1):

    # Open all our logs
    log_files = [h5py.File(log_name) for log_name in log_filenames]

    results = []

    for log_file in log_files:

        # Search for Command items
        for path in log_file.keys():

            m = TASK_PACKET_RE.match(path)
            if m is not None:
                task_packet_group_path = path

                task_packet_path = task_packet_group_path + '/jaiabot.protobuf.TaskPacket'
                task_packets = jaialog_get_object_list(log_file[task_packet_path], repeated_members={"measurement"})

                results += task_packets

    if scheme is not None:
        results = list(filter(lambda object: object['_scheme_']==scheme, results))

    return results

def get_task_packets(log_filenames, scheme=1) -> Iterable[TaskPacket]:
    return [TaskPacket.from_dict(task_packet_json) for task_packet_json in get_task_packet_dicts(log_filenames, scheme)]

def generate_kmz(h5_filename: str, kmz_filename: str):
    task_packets = get_task_packets([h5_filename])
    kmz.write_file(task_packets, kmz_filename)

def getH5File(h5_filename: str):
    '''Returns a Jaia H5 file object'''
    return open(h5_filename, 'br')

# Testing
if __name__ == '__main__':
    filename = '/var/log/jaiabot/bot_offload/bot3_fleet1_20230411T181720.h5'
    generate_kmz(filename, 'test.kmz')


import collections
import heapq
import re

positions = []

for line in """
    Player 1 starting position: 10
    Player 2 starting position: 9
""".strip().splitlines():
    match = re.fullmatch(r'\s*Player \d+ starting position: (\d+)\s*', line)
    positions.append(int(match[1]) - 1)

class DeterministicDiracDie:
    value = 0
    rolls = 0

    def roll(self):
        self.rolls += 1
        self.value %= 100
        self.value += 1
        return self.value

print(positions)

def play(positions, die):
    positions = list(positions)
    scores = [0 for p in positions]
    while True:
        for player_number in range(2):
            for roll_number in range(3):
                positions[player_number] += die.roll()
            positions[player_number] %= 10
            scores[player_number] += positions[player_number] + 1
            if scores[player_number] >= 1000:
                 return player_number, positions, scores, die.rolls

winner, new_positions, scores, rolls = play(positions, DeterministicDiracDie())
print(f'{winner} wins, positions: {new_positions}, scores: {scores}, rolls: {rolls}')

print('Part 1:', rolls * scores[1-winner])

three_roll_outcome_counts = tuple(collections.Counter(
    a + b + c
    for a in range(1, 4)
    for b in range(1, 4)
    for c in range(1, 4)
).items())
print(three_roll_outcome_counts)

universes = [(0, 0, 0, *positions, 0)]
universe_counts = {(0, 0, *positions, 0): 1}
wins = [0, 0]

print_counter = 0
while universes:
    current_entry = heapq.heappop(universes)
    total_score = current_entry[0]
    current_universe = (*current_entry[1:], )
    count = universe_counts.pop(current_universe)
    #print('<-', current_universe, count)
    if print_counter % 2021 == 0:
        print(
            total_score,
            current_universe,
            [len(str(w)) for w in wins],
            len(universes),
            count,
        )
        print_counter = 0
    print_counter += 1
    for roll_sum, roll_count in three_roll_outcome_counts:
        new_universe = list(current_universe)
        player_number = new_universe[-1]
        new_universe[player_number+2] += roll_sum
        new_universe[player_number+2] %= 10
        new_universe[player_number] += new_universe[player_number+2] + 1
        new_universe[-1] = 1 - player_number
        if new_universe[player_number] >= 21:
            wins[player_number] += count * roll_count
            #print('!!')
        else:
            new_tuple = tuple(new_universe)
            new_count = count * roll_count
            if new_tuple in universe_counts:
                universe_counts[new_tuple] += new_count
                #print('++', new_tuple, '*', new_count)
            else:
                universe_counts[new_tuple] = new_count
                total_score = new_universe[0] + new_universe[1]
                heapq.heappush(universes, (total_score, *new_tuple))
                #print('->', new_tuple, '*', new_count)

print(wins)
print('Part 2:', max(wins))


# -*- coding: utf-8 -*-
"""
@author: Akash Meshram (16MT10005)

MIES(EC60091) Assignment
Computer Assignment on Clustering
"""
import pandas as pd
import numpy as np
import operator
import matplotlib.pyplot as plt

#import dataset as numpy array
df = pd.read_excel('dataset.xlsx', header=None).values

f = open('result.txt', 'w')

#setup
k_val = 5 # range of k-centroids
err = [] #list of error for each k value
itern = 20 #no. of iteration on each k value
print('k', 'error')
f.write('k'+" "+ 'error' +"\n")

# iterate over each k
for k in range(1,k_val+1):
	#initialize the centroids, the random 'k' elements in the dataset will be initial centroids
    clusters = df[np.random.randint(df.shape[0], size=k), :]

    #List for bucket for each centroid, where nearest data point will be added
    clusters_group = []
    #Create k bucket for k centroids resp.
    for _ in range(k):
        clusters_group.append([])

    #begin iterations
    for _ in range(itern):
    	#for each point in df
        for row in df:
        	#array for diststace from each centroid
            dist = []
            for c in clusters:
            	#the distance between the point and cluster centers
                dist.append(np.linalg.norm(row-c))
            #choose the nearest centroid, using index
            min_index, min_value = min(enumerate(dist), key=operator.itemgetter(1))
            #adding point to bucket of nearest centroid
            clusters_group[min_index].append(row)

        #average the cluster datapoints to re-calculate the centroids
        for i in range(k):
            temp = np.asarray(clusters_group[i])
            x=np.sum(temp[:,0])/len(temp)
            y=np.sum(temp[:,1])/len(temp)
            clusters[i] = [x,y]
            #for each itration buckets are emptied,
            #except for last, so as to calculate error later
            if (_ != itern-1):
                clusters_group[i] = []

    #claculate error
    error = 0
    for i in range(k):
        temp = np.asarray(clusters_group[i])
        error += np.sum(np.linalg.norm(temp-clusters[i]))

    print(k, error)
    f.write(str(k)+" "+str(error)+'\n')
    #append error to list, so as to plot later
    err.append(error)

#begin plot
fig = plt.figure()
plt.xlabel("K-value")
plt.ylabel("error")
plt.plot(range(1,k_val+1),err,'-r')
fig.savefig('result_plot.png')
f.close()


import json
import time

class Shorting():
    """
    This class is inhereted by TradingBot.
    This class contains all functions for shorting.
    """
    def __init__(self) -> None:
        pass

    def check_if_should_sell_short(self):
        """
        Method for checking if short position should be closed
        """
        current_price = float(self.get_current_value())
        if current_price <= self.target_price:
            self.set_sell_short_order(current_price)
        elif current_price >= self.stop_loss_value:
            self.set_sell_short_order(current_price)

    def set_sell_short_order(self, price):
        """
        Method for creating api call data and resetting shorting variables.
        """
        while True:
            try:
                volume = self.shorted_amount
                api_call_data= {
                    "symbol" : "BTCBUSD",
                    "side" : "BUY",
                    "type" : "LIMIT",
                    "quantity" : self.truncate(volume, 5),
                    "price" : price,
                    "sideEffectType" : "AUTO_REPAY",
                    "timeInForce" : "GTC"
                    }
                self.short_sell_order_api_call(api_call_data)
                data = None
                with open("trades.json", "r", encoding="utf-8") as f:
                    data = json.loads(f.read())
                data["trades"].append(api_call_data)
                with open("trades.json", "w", encoding="utf-8") as f:
                    f.write(json.dumps(data, indent=4))
                print("SELL WAS MADE!")
                self.stop_loss_value = 0
                self.target_price = 0
                self.shorted_amount = 0
                self.loan_open = False
                time.sleep(1)
                break
            except Exception as e:
                time.sleep(1)
                print(f"{e}")
                pass

    def short_sell_order_api_call(self, data):
        """
        This is a separate method for handling closing of short position via the api.
        """
        api_callback_sell = None
        while True:
            try:
                api_callback_sell = self.api.create_margin_order(**data)
                closed_order = self.wait_for_order_to_be_filled(api_callback_sell["clientOrderId"])
                if closed_order == True:
                    return closed_order
            except Exception as e:
                print(data)
                time.sleep(1)
                print(f"{e}")
                pass

    def check_if_should_short(self, averages, data_from_latest_candle, data):
        """
        Method for checking price data in downtrend and searching for possible short openings.
        """
        trend_strength = round(averages[0]["ema9"] / averages[0]["wma30"], 4)
        trend_is_strong = False
        if trend_strength <= 0.997:
            trend_is_strong = True
        current_price = 0
        if float(data_from_latest_candle[4]) < averages[1]["wma30"] and float(data_from_latest_candle[4]) > averages[1]["ema9"]:
            self.short_buy_price = float(data_from_latest_candle[3])
            self.stop_loss_value = round(averages[0]["wma30"] * 1.005, 1)
            print("Last complete candle closed in between indicators SHORT")
            current_price = float(self.get_current_value())
            if current_price <= self.short_buy_price and trend_is_strong:
                print(f"Setting short on {data_from_latest_candle[3]}")
                print(f"With STOP LOSS of {averages[0]['wma30']}")
                volume = 0
                while True:
                    try:
                        max_loan = self.api.get_max_margin_loan(asset='BTC')
                        volume = self.truncate(float(max_loan["amount"]) * 0.5, 5)
                        break
                    except Exception:
                        time.sleep(1)
                        pass
                self.set_short_order(current_price, volume)
            else:
                print(f"current price {current_price}")
                print(f"trend is strong = {trend_is_strong} strength =  {trend_strength}")
        elif self.stop_loss_value != 0 and self.short_buy_price != 0 and trend_is_strong:
            current_price = float(self.get_current_value())
            if current_price <= self.short_buy_price:
                print(f"Setting short on {current_price}")
                print(f"With STOP LOSS of {self.stop_loss_value}")
                volume = 0
                while True:
                    try:
                        max_loan = self.api.get_max_margin_loan(asset='BTC')
                        volume = self.truncate(float(max_loan["amount"]) * 0.5, 5)
                        break
                    except Exception:
                        time.sleep(1)
                        pass
                self.set_short_order(current_price, volume)
        else:
            print(f"current price was not set")
            print(f"trend is strong = {trend_is_strong} strength =  {trend_strength}")

    def set_short_order(self, price, volume):
        while True:
            try:
                print("Trying to short 1")
                api_call_data= {
                    "symbol" : "BTCBUSD",
                    "side" : "SELL",
                    "type" : "LIMIT",
                    "quantity" : volume,
                    "price" : price,
                    "sideEffectType" : "MARGIN_BUY",
                    "timeInForce" : "GTC"
                    }
                data = None
                with open("trades.json", "r", encoding="utf-8") as f:
                    data = json.loads(f.read())
                data["trades"].append(api_call_data)
                with open("trades.json", "w", encoding="utf-8") as f:
                    f.write(json.dumps(data, indent=4))
                self.short_order_api_call(api_call_data)
                self.target_price = round(price * self.short_win_target, 1)
                self.shorted_amount = self.get_borrowed_btc_margin()
                self.loan_open = True
                print(f"Short order set. Target price : {self.target_price}")
                return
            except Exception as e:
                time.sleep(1)
                print(f"{e}")
                pass

    def short_order_api_call(self, data):
        """
        This is a separate method for handling buying via the api.
        """
        while True:
            try:
                api_callback_buy = self.api.create_margin_order(**data)
                print("API CALL BACK WAS RECEIVED")
                closed_order = self.wait_for_order_to_be_filled(api_callback_buy["clientOrderId"])
                if closed_order == True:
                    return closed_order
            except Exception as e:
                time.sleep(1)
                print(f"{e}")
                pass


import random

logo = '''
 _
| |
| |__   __ _ _ __   __ _ _ __ ___   __ _ _ __
| '_ \ / _` | '_ \ / _` | '_ ` _ \ / _` | '_ \
| | | | (_| | | | | (_| | | | | | | (_| | | | |
|_| |_|\__,_|_| |_|\__, |_| |_| |_|\__,_|_| |_|
                    __/ |
                   |___/    '''
print(logo)

stages = ['''
  +---+
  |   |
  O   |
 /|\  |
 / \  |
      |
=========
''', '''
  +---+
  |   |
  O   |
 /|\  |
 /    |
      |
=========
''', '''
  +---+
  |   |
  O   |
 /|\  |
      |
      |
=========
''', '''
  +---+
  |   |
  O   |
 /|   |
      |
      |
=========''', '''
  +---+
  |   |
  O   |
  |   |
      |
      |
=========
''', '''
  +---+
  |   |
  O   |
      |
      |
      |
=========
''', '''
  +---+
  |   |
      |
      |
      |
      |
=========
''']

from hangman_words import word_list

chosen_word = random.choice(word_list)
# print(chosen_word)

life = 6
end_game = False

# creating blanks
display = []
for _ in range(len(chosen_word)):
    display += "_"
# print(display)

while not end_game:
    guess = (input("Guess a letter = ")).lower()
    # print(guess)

    if guess in display:
        print(f"you have already guessed {guess}")

    # checking guessed letter
    for position in range(len(chosen_word)):
        letter = chosen_word[position]
        if letter == guess:
            display[position] = letter
    # print(display)

    if guess not in chosen_word:
        life -= 1
        if life == 0:
            end_game = True
            print("You Lose!!")

    print(f"{' '.join(display)}")

    if "_" not in display:
        end_game = True
        print("You Win!!")

    print(stages[life])


#########################
#    Network Scanner    #
#    By: Yuval Cohen    #
#    Date: 4.4.20       #
#########################

import socket
import threading
from scapy.all import *

def check_ip(ip):
    """
    The function builds a ARP packet and send it to the ip.
    If the packet received successfully it will be added to the ip list.
    :param ip: The ip address.
    """
    global ip_lst

    request = Ether(dst='ff:ff:ff:ff:ff:ff') / ARP(pdst=ip)
    response, _ = srp(request, timeout=1, retry=2, verbose=False)
    results = ()
    for _, packet in response:
        results = (packet.psrc, packet.hwsrc)

    if results:
        ip_lst.append(results)

def main():
    global ip_lst
    local_ip = socket.gethostbyname(socket.gethostname())
    local_ip = local_ip[:local_ip.rfind(".")]

    ip_lst = []
    print("-" * 35)
    print("      Starting the scan...")
    print("-" * 35 + "\n")

    for i in range(255):
        threading.Thread(target=check_ip, args=(f"{local_ip}.{i}",)).start()

    while threading.activeCount() > 1:
        pass

    for addr in ip_lst:
        print(f"Ip address: {addr[0]} | MAC address: {addr[1]}")

    input("\n\nPress any key to continue...")

if __name__ == '__main__':
    main()


import mxnet as mx
from mxnet import nd
from mxnet.gluon import nn

"""
layer = nn.Dense(2)
layer.initialize()#ctx = mx.gpu(0))
x = nd.random_uniform(-1, 1, (3, 4))
layer(x)
print(layer.weight.data)
"""

net = nn.Sequential()
net.add(
    nn.Conv2D(channels=6, kernel_size=5, activation='relu'),
    nn.MaxPool2D(pool_size=2, strides=2),
    nn.Conv2D(channels=16, kernel_size=3, activation='relu'),
    nn.MaxPool2D(pool_size=2, strides=2),
    nn.Dense(120, activation='relu'),
    nn.Dense(84, activation='relu'),
    nn.Dense(10)
)

#print(net)
net.initialize()

x = nd.random.uniform(shape=(4, 1, 28, 28))
y = net(x)
print(y.shape)


class solution:
    def majorityElement(self, nums: list[int]):
        res, cnt = 0, 0
        for n in nums:
            if cnt == 0:
                res = n
            cnt += (1 if res == n else -1)
        return res

x = solution()
print(x.majorityElement(nums=[3, 2, 3]))


from tkinter import *
from tkinter import messagebox
from pkg.fx import *
from decimal import *
from pkg.integral_estimator import *

class IntegralEstimatorGUI:

    def run(self):
        self.root = Tk()
        self.root.title("Integral Estimator")
        self.root.bind('<Return>', lambda event: self.calculate())
        open("geometry.config", "a").close()
        with open("geometry.config", "r") as conf:
            self.root.geometry(conf.read())

        self.root.resizable(False, False)

        self.interface()

        self.root.protocol("WM_DELETE_WINDOW", self.on_close)
        self.root.mainloop()

    def interface(self):
        master = Frame(self.root)

        welcome = Label(master, text="Welcome to Integral Estimator!")
        welcome.pack()

        rules = Frame(master)

        self.rule = StringVar(value="T")
        rule_radio1 = Radiobutton(rules, text="Trapezoidal Rule", variable=self.rule, value="T")
        rule_radio2 = Radiobutton(rules, text="Simpson's Rule", variable=self.rule, value="S")
        rule_radio1.grid(row=0, column=0)
        rule_radio2.grid(row=0, column=1)

        rules.pack()

        entries = Frame(master)

        self.fx = StringVar()
        fx_label = Label(entries, text="f(x) = ")
        fx_label.grid(row=2, column=0)
        fx_entry = Entry(entries, textvariable=self.fx)
        fx_entry.grid(row=2, column=1)
        fx_entry.focus_set()

        self.a = StringVar()
        a_label = Label(entries, text="a = ")
        a_label.grid(row=3, column=0)
        a_entry = Entry(entries, textvariable=self.a)
        a_entry.grid(row=3, column=1)

        self.b = StringVar()
        b_label = Label(entries, text="b = ")
        b_label.grid(row=4, column=0)
        b_entry = Entry(entries, textvariable=self.b)
        b_entry.grid(row=4, column=1)

        self.n = StringVar()
        n_label = Label(entries, text="n = ")
        n_label.grid(row=5, column=0)
        n_entry = Entry(entries, textvariable=self.n)
        n_entry.grid(row=5, column=1)

        entries.pack()

        calc_button = Button(master, text="Compute", command=self.calculate)
        calc_button.pack()

        master.pack(padx=20, pady=20)

    def calculate(self):
        fx = Fx(self.fx.get())
        if not fx.is_valid_fx():
            messagebox.showwarning(message="The funtion of x you entered is invalid!")
            self.root.focus_get().focus_force()
            return

        try:
            a = Decimal(self.a.get())
            b = Decimal(self.b.get())
            n = Decimal(self.n.get())
        except InvalidOperation:
            messagebox.showwarning(message="Please enter a decimal number!")
            self.root.focus_get().focus_force()
            return

        if b < a:
            messagebox.showwarning(message="b must not be equal or less than a!")
            self.root.focus_get().focus_force()
            return

        if self.rule.get() == "S" and (n % 2 != 0 or n <= 0):
            messagebox.showwarning(message="n has to be a positive even number!")
            self.root.focus_get().focus_force()
            return
        elif n <= 0:
            messagebox.showwarning(message="n has to be a positive number!")
            self.root.focus_get().focus_force()
            return

        fxs = []
        for i in range(int(n) + 1):
            x = a + (b - a) / n * i
            fxs.append(fx.evaluate_fx(x))

        integration = IntegralEstimator(self.rule.get(), a, b, n, fxs)

        if integration.rule == "T":
            result = sympy_number_to_str(integration.trap_estimate())
            message = f"The Trapezoidal Rule estimation of the integral of {fx.get_fx()} from {a} to {b} is {result}"
            messagebox.showinfo(message=message)
            self.root.focus_get().focus_force()
        elif integration.rule == "S":
            result = sympy_number_to_str(integration.simp_estimate())
            message = f"The Simpson's Rule estimation of the integral of {fx.get_fx()} from {a} to {b} is {result}"
            messagebox.showinfo(message=message)
            self.root.focus_get().focus_force()

    def on_close(self):
        with open("geometry.config", "w") as conf:
            conf.write(f"+{self.root.winfo_x()}+{self.root.winfo_y()}")
        self.root.destroy()

IntegralEstimatorGUI().run()


#EXERCÍCIO 01 ->    DICIONÁRIO

"""aluno = {}
aluno['nome'] = str(input("Digite o nome do aluno: "))
aluno['media'] = float(input("Digite a média do aluno: "))
if aluno['media'] >= 7 :
    aluno['situação'] = 'Aprovado'
else:
    aluno['situação'] = 'Reprovado'
for k,v in aluno.items():
    print(f'{k} é igual a {v}')"""""

#EXERCÍCIO 02 -> JOGO DE DADOS

"""from random import randint
from time import sleep
from operator import itemgetter
jogo = {
    'jogador1' : randint(1,6),
    'jogador2' : randint(1,6),
    'jogador3' : randint(1,6),
    'jogador4' : randint(1,6),}
ranking = []
print('Valores sorteados: ')
for k,v in jogo.items():
    print(f'{k} tirou {v} no dado.')
    sleep(1)
ranking = sorted(jogo.items(), key=itemgetter(1), reverse= True)
print('-=' * 30)
print('== RANKING DOS JOGADORES ==')
print(ranking)
for i, v in enumerate(ranking):
    print(f'{i+1}º lugar {v[0]} com {v[1]}.')
    sleep(1)"""

#EXERCÍCIO 03 -> CADASTRANDO USUÁRIO EM UM DICIONÁRIO COM UMA PEQUENA CONDIÇÃO

"""from datetime import datetime

cadastro = {}

cadastro['Nome'] = str(input('Digite seu nome: '))
cadastro['idade '] = int(input('Digite o ano em que nasceu: '))
idade = 2022 - cadastro['idade ']
cadastro['idade '] = idade
cadastro['ctps '] = int(input('Digite o número da sua carteira de trabalho (Caso não tenha, digite 0): '))
if cadastro['ctps '] == 0:
    cadastro.update({'ctps': 'Não possui carteira de trabalho.'})
    print('-=' * 35)
    for k,v in cadastro.items():
        print(f'-{k} tem o valor de {v}')

else:
    cadastro['Ano de contratação'] = int(input('Qual foi o ano da sua contratação: '))
    cadastro['Salário '] = int(input('Qual é o salário que você recebe: '))
    cadastro['aposentadoria'] = ((cadastro['Ano de contratação'] + 35) - datetime.now().year)
    print('-=' * 35)
    for k,v in cadastro.items():
        print(f'-{k} tem o valor de {v}')"""

# EXERCÍCIO 04 -> CADASTRO DE JOGADOR DE FUTEBOL

"""jogador = {}
gols = []
jogador['nome'] = str(input('Qual o nome do jogador?: '))
partidas = int(input(f'Quantas partidas o {jogador["nome"]} jogou?: '))
for i in range(0, partidas):
    cont = 1
    cont = cont + i
    partida = int(input(f'Quantos gols na {cont}ª partida? : '))
    gols.append(partida)
    partida += partida
    jogador['gols'] = gols
    jogador['Total'] = partida

print(jogador)
print('-='*30)
for k,v in jogador.items():
    print(f'O campo {k} tem o valor de {v}')
print('-='*30)
print(f'O jogador {jogador["nome"]} jogou {partidas} partidas.')
for i, v in enumerate(jogador['gols']):
    print(f' => Na partida {i}, marcou {v} gols. ')
print(f'Foi um total de {jogador["Total"]} gols.')"""

# EXERCÍCIO 05 -> Unindo Dicionários e Listas (Cadastro e Análise de Dados)
"""from ast import While

galera = []
pessoa = {}
soma = média = 0
while True:
    pessoa.clear()
    pessoa['nome'] = str(input('Nome: '))
    while True:
        pessoa['sexo'] = str(input('Sexo: [M/F] ')).upper()[0]
        if pessoa['sexo'] in 'MF':
            break
        print('ERRO! Por favor digite, apenas M ou F ')
    pessoa['idade'] = int(input('Idade: '))
    soma += pessoa['idade']
    galera.append(pessoa.copy())
    while True:
        resp = str(input('Quer continuar? [S/N] ')).upper()[0]
        if resp in 'SN':
            break
        print('ERRO! Responda apenas Sim ou Não')
    if resp == 'N':
        break
print('-=' * 30)
print(galera)
print(f'A) Ao todo temos {len(galera)} pessoas cadastradas.')
média = soma / len(galera)
print(f'B) A média de idade é de {média:5.2f} anos.')
print(f'C) As mulheres cadastradas foram ', end=' ')
for p in galera:
    if p['sexo'] == 'F':
        print(f'{p["nome"]} ', end='')
print()
print('D) Listas das pessoas que estão acima da média: ', end='')
for p in galera:
    if p['idade'] >= média:
        print('  ')
        for k,v in p.items():
            print(f'{k } = {v}', end=' ')
        print()
print("<<ENCERRADO>>")"""

#EXERCÍCIO 06 -> CADASTRANDO JOGADORES E ADICIONANDO EM UMA LISTA

time = []
jogador = {}
partidas = []

while True:
    jogador.clear()
    jogador['nome'] = str(input('Digite o nome do Jogador: '))
    tot = int(input(f'Quantas partidas {jogador["nome"]} jogou? '))
    partidas.clear()
    for c in range (0, tot):
        partidas.append(int(input(f'Quantos gol na partida {c+1}? ')))
    jogador['gols'] = partidas[:]
    jogador['total'] = sum(partidas)
    time.append(jogador.copy())
    while True:
        resp = str(input('Quer continuar? [S/N] ')).upper()[0]
        if resp in 'SN':
            break
        print('ERRO! RESPONDA APENAS S OU N.')
    if resp == 'N':
        break
print('-='* 30)
print('cod ', end='')
for i in jogador.keys():
    print(f'{i:<15}', end='')
print()
print('-=' * 40 )
for k,v in enumerate(time):
    print(f'{k:>3} ', end='')
    for d in v.values():
        print(f'{str(d):<15}', end='')
    print()
print('-'*40 )
while True:
    busca = int(input('Mostrar dados de qual jogador? (999 para parar )'))
    if busca == 999:
        break
    if busca >= len(time):
        print(f'Não existe jogador com código {busca}!')
    else:
        print(f' -- LEVANTAMENTO DO JOGADOR {time[busca]["nome"]}')
        for i,g in enumerate(time[busca]['gols']):
            print(f'No jogo {i+1} fez {g} gols')
        print('-' * 40)
print('<<< VOLTE SEMPRE >>>')


import scrapy
from scrapy import FormRequest
from ..items import LawyerSearchItem
from pyexcel_xlsx import get_data

class LawyerSpider(scrapy.Spider):
    name = "lawyer"

    allowed_domains = ["texasbar.com"]

    start_urls = [
        "https://www.texasbar.com/AM/Template.cfm?Section=Find_A_Lawyer&Template=/CustomSource/MemberDirectory/Search_Form_Client_Main.cfm"]

    lawyers = get_data('./Names.xlsx')
    lawyers_list = list(lawyers.values())[0]
    fixed_list = lawyers_list[1:len(lawyers_list)-1]

    def parse(self, response):
        for person in self.fixed_list:
            yield FormRequest.from_response(
                response,
                formname="SearchForm_Public",
                url="/AM/Template.cfm?Section=Find_A_Lawyer&Template=/CustomSource/MemberDirectory/Result_form_client.cfm",
                formdata={
                    "FirstName": person[1],
                    "LastName": person[0],
                },
                callback=self.parse_result,
                dont_filter=True,
            )
    def parse_result(self, response):
        item = LawyerSearchItem()
        try:
            city, rest = response.xpath('//div[@class="avatar-column"]/p[@class="address"]/text()')[1].extract().strip().split(',')
            state, zip_num = rest.split()
        except IndexError:
            city = 'No info'
            state = 'No info'
            zip_num = 'No info'
        try:
            item['firm_name'] = response.xpath('//div[@class="avatar-column"]/h5/text()')[0].extract().strip()
        except IndexError:
            item['firm_name'] = 'No info'
        try:
            item['phone'] = response.xpath('//div[@class="contact"]/a[2]/text()')[0].extract().strip().split(':')[1].strip()
        except IndexError:
            item['phone'] = 'No info'
        try:
            item['name'] = response.xpath('//div[@class="avatar-column"]/h3/a/span[@class="given-name"][1]/text()')[0].extract().strip()
        except IndexError:
            item['name'] = 'No info'
        try:
            item['address'] = response.xpath('//div[@class="avatar-column"]/p[@class="address"]/text()')[0].extract().strip()
        except IndexError:
            item['address'] = 'No info'
        item['city'] = city
        item['state'] = state
        item['zip_num'] = zip_num
        item['areas'] = response.xpath('//div[@class="avatar-column"]/p[@class="areas"]/text()')[1].extract().strip()
        yield item


#1시간 #그래프이론 #그래프탐색 #브루트포스알고리즘 #BFS #DFS
"""
https://www.acmicpc.net/problem/2468
[2468] : 안전 영역
"""

import sys
input = sys.stdin.readline
sys.setrecursionlimit(10**6)

N = int(input())
arr = [list(map(int, input().split())) for _ in range(N)]
max_height = max(map(max, arr))
min_height = min(map(min, arr))
dx = [0, 0, -1, 1]
dy = [-1, 1, 0, 0]

def dfs(x, y, height):
    visited[x][y] = 1

    for i in range(4):
        a = x + dx[i]
        b = y + dy[i]

        if 0<=a<N and 0<=b<N and arr[a][b] > height and visited[a][b] == 0:
            visited[a][b] = 1
            dfs(a, b, height)

res = 0
for height in range(min_height-1, max_height+1):
    cnt = 0
    visited = [[0] * N for _ in range(N)]
    for i in range(N):
        for j in range(N):
            if arr[i][j] > height and visited[i][j] == 0:
                dfs(i, j, height)
                cnt += 1
    if cnt > res:
        res = cnt

print(res)


from datetime import datetime
from typing import List, Literal, Optional, Tuple
from uuid import UUID

from pybotx.client.authorized_botx_method import AuthorizedBotXMethod
from pybotx.missing import Missing, Undefined
from pybotx.models.api_base import UnverifiedPayloadBaseModel, VerifiedPayloadBaseModel
from pybotx.models.bot_catalog import BotsListItem

class BotXAPIBotsListRequestPayload(UnverifiedPayloadBaseModel):
    since: Missing[datetime] = Undefined

    @classmethod
    def from_domain(
        cls,
        since: Missing[datetime] = Undefined,
    ) -> "BotXAPIBotsListRequestPayload":
        return cls(since=since)

class BotXAPIBotItem(VerifiedPayloadBaseModel):
    user_huid: UUID
    name: str
    description: str
    avatar: Optional[str] = None
    enabled: bool

class BotXAPIBotsListResult(VerifiedPayloadBaseModel):
    generated_at: datetime
    bots: List[BotXAPIBotItem]

class BotXAPIBotsListResponsePayload(VerifiedPayloadBaseModel):
    result: BotXAPIBotsListResult
    status: Literal["ok"]

    def to_domain(self) -> Tuple[List[BotsListItem], datetime]:
        bots_list = [
            BotsListItem(
                id=bot.user_huid,
                name=bot.name,
                description=bot.description,
                avatar=bot.avatar,
                enabled=bot.enabled,
            )
            for bot in self.result.bots
        ]
        return bots_list, self.result.generated_at

class BotsListMethod(AuthorizedBotXMethod):
    async def execute(
        self,
        payload: BotXAPIBotsListRequestPayload,
    ) -> BotXAPIBotsListResponsePayload:
        path = "/api/v1/botx/bots/catalog"

        response = await self._botx_method_call(
            "GET",
            self._build_url(path),
            params=payload.jsonable_dict(),
        )

        return self._verify_and_extract_api_model(
            BotXAPIBotsListResponsePayload,
            response,
        )


from Crypto import Random
from Crypto.Cipher import AES
from fantastico.oauth2.exceptions import OAuth2InvalidTokenDescriptorError, OAuth2TokenEncryptionError
from fantastico.oauth2.token import Token
from fantastico.oauth2.token_encryption import AesTokenEncryption
from fantastico.tests.base_case import FantasticoUnitTestsCase

class AesTokenEncryptionTests(FantasticoUnitTestsCase):
    '''This class provides tests suite for AES token encryption provider from fantastico.'''

    _aes_iv = None

    def init(self):
        '''This method is invoked automatically in order to initialize common test cases dependencies.'''

        self._aes_iv = Random.new().read(AES.block_size)

    def test_aes_ok(self):
        '''This test case ensures token encryption correctly supports AES-128 algorithm.'''

        aes_provider = AesTokenEncryption()

        for key_length in [16, 24, 32]:
            token_key = Random.new().read(key_length)

            token = Token({"client_id": "test client",
                           "attr1": "test"})

            token_encrypted = aes_provider.encrypt_token(token, self._aes_iv, token_key)
            token_decrypted = aes_provider.decrypt_token(token_encrypted, self._aes_iv, token_key)

            self.assertIsNotNone(token_decrypted)
            self.assertEqual(token.client_id, token_decrypted.client_id)
            self.assertEqual(token.attr1, token_decrypted.attr1)

    def test_aes_encrypt_notoken(self):
        '''This test case ensures aes encrypt fails if no token is given.'''

        aes_provider = AesTokenEncryption()

        with self.assertRaises(OAuth2InvalidTokenDescriptorError) as ctx:
            aes_provider.encrypt_token(None, None, None)

        self.assertEqual("token", ctx.exception.attr_name)

    def test_aes_encrypt_notokeniv(self):
        '''This test case ensures aes encrypt fails if no token_iv is given.'''

        aes_provider = AesTokenEncryption()

        with self.assertRaises(OAuth2InvalidTokenDescriptorError) as ctx:
            aes_provider.encrypt_token(Token({}), None, None)

        self.assertEqual("token_iv", ctx.exception.attr_name)

    def test_aes_encrypt_notokenkey(self):
        '''This test case ensures aes encrypt fails if no token_key is given.'''

        aes_provider = AesTokenEncryption()

        with self.assertRaises(OAuth2InvalidTokenDescriptorError) as ctx:
            aes_provider.encrypt_token(Token({}), "simple key", None)

        self.assertEqual("token_key", ctx.exception.attr_name)

    def test_aes_encrypt_unexpected_error(self):
        '''This test case ensures all unexpected exceptions from encrypt are converted correctly into OAuth2 encryption errors.'''

        aes_provider = AesTokenEncryption()

        with self.assertRaises(OAuth2TokenEncryptionError):
            aes_provider.encrypt_token(Token({}), "Simple IV", "Simple Key")

    def test_aes_decrypt_notoken(self):
        '''This test case ensures aes decrypt fails if no token is given.'''

        aes_provider = AesTokenEncryption()

        with self.assertRaises(OAuth2InvalidTokenDescriptorError) as ctx:
            aes_provider.decrypt_token(None, None, None)

        self.assertEqual("encrypted_str", ctx.exception.attr_name)

    def test_aes_descrypt_notokeniv(self):
        '''This test case ensures aes decrypt fails if no token iv is given.'''

        aes_provider = AesTokenEncryption()

        with self.assertRaises(OAuth2InvalidTokenDescriptorError) as ctx:
            aes_provider.decrypt_token(Token({}), None, None)

        self.assertEqual("token_iv", ctx.exception.attr_name)

    def test_aes_decrypt_notokenkey(self):
        '''This test case ensures aes decrypt fails if no token_key is given.'''

        aes_provider = AesTokenEncryption()

        with self.assertRaises(OAuth2InvalidTokenDescriptorError) as ctx:
            aes_provider.decrypt_token(Token({}), "simple key", None)

        self.assertEqual("token_key", ctx.exception.attr_name)

    def test_aes_decrypt_unexpected_error(self):
        '''This test case ensures all unexpected exceptions from decrypt are converted correctly into OAuth2 encryption errors.'''

        aes_provider = AesTokenEncryption()

        with self.assertRaises(OAuth2TokenEncryptionError):
            aes_provider.decrypt_token(Token({}), "Simple IV", "Simple Key")


import requests

def obter_localizacao_por_ip():
    url = 'https://ipinfo.io/json'
    response = requests.get(url)
    data = response.json()
    cidade = data.get('city', '')
    return cidade

def obter_clima(cidade):
    api_key = 'SUA_CHAVE_DE_API'
    url = f'http://api.openweathermap.org/data/2.5/weather?q={cidade}&appid={api_key}'

    response = requests.get(url)
    dados = response.json()

    if dados['cod'] == 200:
        descricao_clima = dados['weather'][0]['description']
        temperatura = dados['main']['temp']
        umidade = dados['main']['humidity']

        print(f"Clima em {cidade}: {descricao_clima}")
        print(f"Temperatura: {temperatura} K")
        print(f"Umidade: {umidade}%")
    else:
        print(f"Erro: {dados['message']}")


"""
* assert (verificação) é como se fosse um if
porém é usado para teste

* Podemos usar vários asserts para indicarmos
diversas verificações de erros

* E também podemos colocar uma mensagem
de acordo com o contexto do erro
"""

from modulo1 import somar, calcular_fatorial

def teste1():
    try:
        resultado = somar(10, 20)
        assert resultado > 0, 'Resultado não é positivo'
        assert resultado % 2 == 0, 'Resultado não é par'
        assert resultado == 30, 'Resultado não é 30'
        print('CORRETO')
    except AssertionError as x:
        print('ERRO:', resultado, x)

def teste2():
    try:
        resultado = somar(-50, -50)
        assert resultado == -100
        print('CORRETO')
    except AssertionError:
        print('ERRO')

def teste_fatorial():
    try:
        resultado = calcular_fatorial(5)
        assert resultado == 120
        print('Fatorial CORRETO')
    except AssertionError:
        print('ERRO no fatorial')

teste1()
teste2()
teste_fatorial()


# -*- coding: utf-8 -*-
"""
Created on Thu Oct 22 14:03:16 2020

@author: 張濟
"""

import csv
import os
import shutil

from torchvision import models
import torch
import torch.nn as nn
import torch.utils.data as Data
from torchvision.datasets import ImageFolder
import torchvision
from torch.optim import lr_scheduler
from PIL import Image

import time
#GPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'
#Hyperparametes
BATCH_SIZE = 64
LR = 0.1  #0.05 for resnet
EPOCH = 200

def train_model(model, criterion, optimizer, scheduler, num_epochs, train_loader, valid_loader):
    model = model.train()
    for epoch in range(num_epochs):
        training_corrects = 0
        valid_corrects = 0
        for step, (bx, by) in enumerate(train_loader):
            bx, by = bx.to(device), by.to(device)
            output = model(bx)
            _, preds = torch.max(output, 1)
            loss = criterion(output, by)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            training_corrects += torch.sum(preds == by.data)

        with torch.no_grad():
            model = model.eval()
            for (inputs, labels) in valid_loader:
                inputs, labels = inputs.to(device), labels.to(device)
                output = model(inputs)
                _, preds = torch.max(output, 1)
                valid_corrects += torch.sum(preds == labels.data)
            model = model.train()
        print("train correct : %d"%training_corrects)   #all 8942
        print("validation correct : %d"%valid_corrects) #all 2237
        scheduler.step()
        #print("accuracy : %d"%epoch_acc)
    return model

def predict_carbrand(model, classes, img, imgid, writer):
    imgid = imgid.split('.')[0]
    model = model.eval()
    with torch.no_grad():
        prob = model(img.to(device))
        v,l = torch.max(prob,1)
        pred_r = classes[l.data.item()]
        writer.writerow([imgid, pred_r])
def copy_img(mode, imgclass_index, pathold):
    pathnew = ""
    if mode == "train":
        if imgclass_index<10 and imgclass_index>=0:
            pathnew = "./training_data_p/00" + str(imgclass_index)
        elif imgclass_index<100 and imgclass_index>=10:
            pathnew = "./training_data_p/0" + str(imgclass_index)
        else:
            pathnew = "./training_data_p/%d"%(imgclass_index)
        #if not os.path.exists(pathnew):
        #    shutil.copy(pathold, pathnew)
    else:
        if imgclass_index<10 and imgclass_index>=0:
            pathnew = "./valid_set/00" + str(imgclass_index)
        elif imgclass_index<100 and imgclass_index>=10:
            pathnew = "./valid_set/0" + str(imgclass_index)
        else:
            pathnew = "./valid_set/%d"%(imgclass_index)
        #if not os.path.exists(pathnew):
        #    shutil.copy(pathold, pathnew)
    return pathnew
def main():
    numofimg = 11185
    with open("training_labels.csv", newline="") as csvfile:
        rows = csv.reader(csvfile)
        i = 0
        classes = []
        for row in rows:
            if i == 0:
                i+=1
                continue
            classname = str(row[1])
            if i-1<10 and i-1>=0:
                path_t = "./training_data_p/00" + str(i-1)
                path_v = "./valid_set/00" + str(i-1)
            elif i-1<100 and i-1>=10:
                path_t = "./training_data_p/0" + str(i-1)
                path_v = "./valid_set/0" + str(i-1)
            else:
                path_t = "./training_data_p/%d"%(i-1)
                path_v = "./valid_set/%d"%(i-1)
            if not classname in classes:
                classes.append(classname)
                if not os.path.exists(path_t):
                    os.mkdir(path_t)
                    os.mkdir(path_v)
                    i+=1

    with open("training_labels.csv", newline="") as csvfile:
        i = 0
        rows = csv.reader(csvfile)
        for row in rows:
            if i == 0:
                i+=1
                continue
            imgid = str(row[0])
            pathold = "./training_data/" + imgid + ".jpg"
            imgclass = str(row[1])
            imgclass_index = classes.index(imgclass)
            if i <= numofimg*0.8:
                pathnew = copy_img("train", imgclass_index, pathold)
            else:
                pathnew = copy_img("valid", imgclass_index, pathold)

            if not os.path.exists(pathnew):
                shutil.copy(pathold, pathnew)
            i+=1
    data_transform = {
        # train use data augmentation
        'train':
            torchvision.transforms.Compose([
                torchvision.transforms.Resize((256, 256)),
                torchvision.transforms.RandomRotation(degrees=30),
                torchvision.transforms.RandomCrop(size=(224,224)),
                #(lambda crops: torch.stack([torchvision.transforms.ToTensor()(crop) for crop in crops])),
                #torchvision.transforms.ColorJitter(),
                #torchvision.transforms.RandomHorizontalFlip(),
                #torchvision.transforms.RandomResizedCrop(size=(224, 224), scale=(0.85,1), ratio=(0.9, 1.1)),

                #torchvision.transforms.Resize((224, 224)),
                #torchvision.transforms.RandomHorizontalFlip(),
                torchvision.transforms.ToTensor(),
                torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225])
            ]),
        #validation does not need augmentation
        'valid':
            torchvision.transforms.Compose([
                torchvision.transforms.Resize((224, 224)),
                torchvision.transforms.ToTensor(),
                torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std = [0.229, 0.224, 0.225])
            ])
    }

    train_data = ImageFolder(
        root = os.path.join("training_data_p"),
        transform = data_transform['train']
    )
    train_loader = Data.DataLoader(
        dataset = train_data,
        batch_size = BATCH_SIZE,
        shuffle = True,
        #num_workers = 2
    )
    valid_data = ImageFolder(
        root = os.path.join("valid_set"),
        transform = data_transform['valid']
    )
    valid_loader = Data.DataLoader(
        dataset = valid_data,
        batch_size = BATCH_SIZE,
        shuffle = True,
        #num_workers = 2
    )

    print(type(train_data))

    print("num of training data : %d" %(len(train_data)))
    print("num of validation data : %d" %(len(valid_data)))
    print("epoch num:%d"%EPOCH)
    #print(train_data.class_to_idx)
    model = models.densenet161(pretrained=True)
    #model = models.resnet50(pretrained = True)
    #freeze the parameters so that gradiants can't compute in backward()
    #print(model)
    ct = 0
    for child in model.children():
        ct += 1
        if ct < 20:    #(6)for resnset50
            for param in child.parameters():
                param.requires_grad = False
    """
    for param in model.parameters():
        param.requires_grad = False
    """
    #Parameters of newly constructed modules have requires_grad=True by default
    #num_ftrs = model.fc.in_features
    #model.fc = nn.Linear(num_ftrs, 196)

    #for densenet
    num_ftrs = model.classifier.in_features
    model.classifier = nn.Linear(num_ftrs, 196)

    """
    model.fc = nn.Sequential(
        nn.Linear(num_ftrs, 1024),
        nn.Dropout(0.4),
        nn.ReLU(),
        nn.Linear(1024, 196)
        )
    """
    #print(model)
    model = model.to(device)

    #optimizer = torch.optim.SGD(model.parameters(), lr=LR)
    optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=LR)
    #optimizer = torch.optim.Adam(model.parameters(), lr = LR)
    loss_func = nn.CrossEntropyLoss()
    step_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=25, gamma=0.2)
    ts = time.time()
    model = train_model(model, loss_func, optimizer, step_lr_scheduler, EPOCH, train_loader, valid_loader)
    te = time.time()
    print("time : %d" %(te-ts))
    #test img transform
    """
    img = img = Image.open("./testing_data/000093.jpg")
    img.show()
    img_t = torchvision.transforms.Compose([torchvision.transforms.Resize((256, 256)),
                                            torchvision.transforms.RandomCrop(size=(224,224)),
                                            #(lambda crops: torch.stack([torchvision.transforms.ToTensor()(crop) for crop in crops])),
                                            #torchvision.transforms.ToTensor()
                                            ])
    img = img_t(img)
    print(type(img))
    #print(img.size())
    img.show()
    """

    #several cases
    imgs = os.listdir("./testing_data")
    with open("output.csv","w",newline='') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerow(["id", "label"])
        i = 0
        for img in imgs:
            imgid = img
            img = Image.open("./testing_data/"+img)
            try:
                img = data_transform['valid'](img)
            except:
                print("%s set into RGB format"%imgid)
                img = data_transform['valid'](img.convert('RGB'))
            img = torch.unsqueeze(img,0)
            #print(type(img))
            #print(img.shape)
            predict_carbrand(model, classes, img, imgid, writer)

if __name__ == "__main__":
    main()


import pytest

import merlin.models.tf as ml

tf = pytest.importorskip("tensorflow")

@pytest.fixture
def fifo_queue_fixture():
    queue = ml.FIFOQueue(
        capacity=10,
        dims=[5],
        dtype=tf.float32,
    )

    return queue

def test_queue_single_enqueue_dequeue(fifo_queue_fixture):
    queue = fifo_queue_fixture

    input = tf.random.uniform((5,))
    queue.enqueue(input)

    output = queue.dequeue()
    assert tf.shape(output)[0] == 5
    assert tf.reduce_all(input == output)

def test_queue_multiple_single_enqueue_dequeue(fifo_queue_fixture):
    queue = fifo_queue_fixture

    inputs = [tf.random.uniform((5,)) for _ in range(4)]
    for input in inputs:
        queue.enqueue(input)

    for input in inputs:
        output = queue.dequeue()
        assert tf.reduce_all(input == output)

def test_queue_dequeue_error_when_fully_emptied(fifo_queue_fixture):
    queue = fifo_queue_fixture

    input = tf.random.uniform((5,))
    queue.enqueue(input)

    _ = queue.dequeue()

    with pytest.raises(IndexError) as excinfo:
        _ = queue.dequeue()
    assert "The queue is empty" in str(excinfo.value)

    with pytest.raises(IndexError) as excinfo:
        _ = queue.dequeue_many(3)
    assert "The queue is empty" in str(excinfo.value)

def test_queue_dequeue_error_when_nothing_added(fifo_queue_fixture):
    queue = fifo_queue_fixture

    with pytest.raises(IndexError) as excinfo:
        _ = queue.dequeue()
    assert "The queue is empty" in str(excinfo.value)

    with pytest.raises(IndexError) as excinfo:
        _ = queue.dequeue_many(3)
    assert "The queue is empty" in str(excinfo.value)

def test_queue_enqueue_dequeue_many(fifo_queue_fixture):
    queue = fifo_queue_fixture

    input = tf.random.uniform((3, 5))
    queue.enqueue_many(input)

    output = queue.dequeue_many(3)
    assert tf.shape(output)[0] == 3
    assert tf.shape(output)[1] == 5
    assert tf.reduce_all(input == output)

def test_queue_enqueue_last_dequeue_many_smaller(fifo_queue_fixture):
    queue = fifo_queue_fixture

    outputs_current = queue.list_all()
    assert tf.shape(outputs_current)[0] == 0

    input = tf.random.uniform((3, 5))
    queue.enqueue_many(input)

    outputs_current = queue.list_all()
    assert tf.reduce_all(outputs_current == input)

    output = queue.dequeue_many(2)
    assert tf.shape(output)[0] == 2
    assert tf.reduce_all(output == input[:2])

    outputs_current = queue.list_all()
    assert tf.reduce_all(outputs_current == input[2:])

    output = queue.dequeue_many(2)
    assert tf.shape(output)[0] == 1
    assert tf.reduce_all(output == input[2:])

    outputs_current = queue.list_all()
    assert tf.shape(outputs_current)[0] == 0

def test_queue_enqueue_until_exceeds_capacity(fifo_queue_fixture):
    queue = fifo_queue_fixture

    inputs = list([tf.random.uniform((3, 5)) for _ in range(4)])
    for input in inputs:
        queue.enqueue_many(input)

    outputs_list = queue.list_all()
    # After adding 12 elements to a 10-sized queue, should keep the last 10 added
    assert tf.reduce_all(outputs_list == tf.concat(inputs, axis=0)[2:])

def test_queue_enqueue_dequeue_consistent_size_and_list(fifo_queue_fixture):
    queue = fifo_queue_fixture

    for i in range(1, 12):
        queue.enqueue(tf.random.uniform((5,)))
        if i < queue.capacity:
            assert queue.count() == i
            assert queue.list_all().shape[0] == i
            assert not queue.at_full_capacity
        else:
            assert queue.count() == queue.capacity
            assert queue.list_all().shape[0] == queue.capacity
            assert queue.at_full_capacity

    for i in range(9, 0):
        _ = queue.dequeue()
        assert queue.count() == i
        assert queue.list_all().shape[0] == i
        assert not queue.at_full_capacity

def test_queue_clear(fifo_queue_fixture):
    queue = fifo_queue_fixture

    for i in range(1, 12):
        queue.enqueue(tf.random.uniform((5,)))

    assert queue.count() == queue.capacity
    assert queue.list_all().shape[0] == queue.capacity
    assert queue.at_full_capacity

    queue.clear()
    assert queue.count() == 0
    assert queue.list_all().shape[0] == 0
    assert not queue.at_full_capacity

def test_enqueue_tensors_1_dim_int32():
    queue = ml.FIFOQueue(
        capacity=10,
        dims=[],
        dtype=tf.int32,
    )

    single_input = tf.random.uniform((), minval=1, maxval=10000, dtype=tf.int32)
    queue.enqueue(single_input)

    multiple_inputs = tf.random.uniform((7,), minval=1, maxval=10000, dtype=tf.int32)
    queue.enqueue_many(multiple_inputs)

def test_enqueue_tensors_wrong_dim():
    queue = ml.FIFOQueue(
        capacity=10,
        dims=[13],
        dtype=tf.float32,
    )

    single_input = tf.random.uniform((12,))
    with pytest.raises(AssertionError) as excinfo:
        queue.enqueue(single_input)
    assert "The shape of val and self.dims should match" in str(excinfo.value)

    multiple_inputs = tf.random.uniform((7, 11))
    with pytest.raises(AssertionError) as excinfo:
        queue.enqueue_many(multiple_inputs)
    assert (
        "The shape of values (ignoring the first dim which is the number of examples) "
        "and self.dims should match"
    ) in str(excinfo.value)

def test_indexof():
    queue = ml.FIFOQueue(
        capacity=10,
        dims=[],
        dtype=tf.int32,
    )

    item_ids = tf.range(10, 0, -1, dtype=tf.int32)
    queue.enqueue_many(item_ids)

    indices = queue.index_of([0, 1, 2])
    tf.assert_equal(tf.cast(indices, tf.int32), [-1, 9, 8])

def test_get_values_by_indices():
    queue = ml.FIFOQueue(
        capacity=10,
        dims=[],
        dtype=tf.int32,
    )

    item_ids = tf.range(10, 0, -1, dtype=tf.int32)
    queue.enqueue_many(item_ids)

    values = queue.get_values_by_indices([1, 3])
    tf.assert_equal(values, [9, 7])

def test_update_by_indices():
    queue = ml.FIFOQueue(
        capacity=10,
        dims=[],
        dtype=tf.int32,
    )

    item_ids = tf.range(10, 0, -1, dtype=tf.int32)
    queue.enqueue_many(item_ids)

    queue.update_by_indices(indices=tf.constant([[1], [2]]), values=tf.constant([20, 21]))
    values = queue.list_all()
    tf.assert_equal(values, [10, 20, 21, 7, 6, 5, 4, 3, 2, 1])


#
# @lc app=leetcode id=3 lang=python3
#
# [3] Longest Substring Without Repeating Characters
#

# @lc code=start
class Solution:
    def lengthOfLongestSubstring(self, s: str) -> int:
        answer = 0
        start = 0
        end = 0

        length = len(s)
        hash = set()

        while end < length:

            if s[end] in hash:
                hash.remove(s[start])
                start += 1
            else:
                hash.add(s[end])
                end += 1

                l = end - start
                answer = l if l > answer else answer

        return answer
# @lc code=end

#%%
print(Solution().lengthOfLongestSubstring('pwwkew'))

#%%


import re
from pathlib import Path
import json

REACT_ICONS_BASE_DIR = Path(__file__).parent.parent/"react-icons"

icons={}

def extractSvg(code,dir_name):
    # Define the regular expression pattern to match the function name and JSON object
    pattern = r'(\w+)\s*=\s*function\s*(\w+)\s*\(([^)]*)\)\s*{.*?GenIcon\((\{.*?\})\)\(props\);'

    # Find all occurrences of the pattern in the code
    matches = re.finditer(pattern, code, re.DOTALL)

    # Extract the function names and JSON objects from each match
    d={}
    for match in matches:
        function_name = match.group(2)
        json_object_str = match.group(4)
        d[function_name] = eval(json_object_str)

    if d:
        icons[dir_name] = d

for path in REACT_ICONS_BASE_DIR.iterdir():
    if not path.is_dir():
        continue

    js_path = path/"index.js"
    if not js_path.exists():
        continue

    print(path.stem)

    with open(js_path,"r") as fh:
        extractSvg(fh.read(),path.stem)

# with open(Path(__file__).parent/"svgIcons.json","w") as fh:
#     json.dump(icons,fh,indent=4)

# print("start writing python file")
# ICONS_PYTHON_FOLDER = Path(__file__).parent.parent/"pyIcons/icons"
# for key,values in icons.items():
#     with open(ICONS_PYTHON_FOLDER/f"{key}.py","w") as fh:
#         print(f"writing {key}.py ....")
#         content = ""
#         for iconName,icon in values.items():
#             content +=f"\n{iconName} = {icon}"
#         fh.write(content)

imports = """
from typing import List,Union
from ..utils import create_svg_element
"""
fun = """
def {}(width:str="1em",height:str="1em",stroke:str="none",fill:Union[str,None]="black",stroke_width:str="0",class_names:List[str]=[],id:str=""):
    return create_svg_element({},width,height,stroke,fill,stroke_width,class_names,id)
"""
print("start writing python file")
ICONS_PYTHON_FOLDER = Path(__file__).parent.parent/"pyIcons/icons"
for key,values in icons.items():
    with open(ICONS_PYTHON_FOLDER/f"{key}.py","w") as fh:
        print(f"writing {key}.py ....")
        content = imports
        for iconName,icon in values.items():
            content +=fun.format(iconName,icon)
        fh.write(content.strip())


import os
import time

import matplotlib.pyplot as plt
import numpy as np
import scipy.sparse as sparse
import scipy.sparse.linalg as linalg

from sde_hjb_solver.utils_path import save_data, load_data

class SolverHJB2D(object):
    ''' This class provides a solver of the following 2d BVP by using a
        finite differences method:
            0 = LΨ − f Ψ in S
            Ψ = exp(− g) in ∂S,
        where f and g are the running and terminal costs of the work functional,
        the solution Ψ is the quantity we want to estimate and L is the infinitessimal
        generator of the not controlled 2d diffusion process:
            L = b(x)·∇ + 1/2 sigma(x)^2·Δ

    Attributes
    ----------
    sde: ControlledSDE object
        controlled sde object
    ct_initial: float
        initial computational time
    ct_time: float
        final computational time
    ct: float
        computational time
    psi: array
        solution of the BVP problem
    solved: bool
        flag telling us if the problem is solved
    value_function: array
        value function of the HJB equation
    u_opt: array
        optimal control of the HJB equation

    Methods
    -------
    __init__(sde, h, load)

    start_timer()

    stop_timer()

    get_flatten_index(idx)

    get_bumpy_index(idx)

    get_x(k)

    get_flatten_idx_from_axis_neighbours(idx, i)

    solve_bvp()

    compute_value_function()

    compute_optimal_control()

    save()

    load()

    get_psi_at_x(x)

    get_value_function_at_x(x)

    get_u_opt_at_x(x)

    get_perturbed_potential_and_drift()

    write_report(x)

    plot_2d_psi(ylim=None)

    plot_2d_value_function(ylim=None)

    plot_2d_perturbed_potential(ylim=None)

    plot_2d_control(ylim=None)

    plot_2d_perturbed_drift(ylim=None)

    '''

    def __init__(self, sde, h, load=False):
        ''' init method

        Parameters
        ----------
        sde: langevinSDE object
            overdamped langevin sde object
        h: float
            step size
        load: bool
            load solution

        Raises
        ------
        NotImplementedError
            If dimension d is greater than 1
        '''

        if sde.d != 2:
            raise NotImplementedError('d > 2 not supported')

        # sde object
        self.sde = sde

        # discretization step
        self.h = h

        # rel directory path
        self.rel_dir_path = os.path.join(sde.name, 'h{:.0e}'.format(h))

        if load:
            self.load()

    def start_timer(self):
        ''' start timer
        '''
        self.ct_initial = time.perf_counter()

    def stop_timer(self):
        ''' stop timer
        '''
        self.ct_final = time.perf_counter()
        self.ct = self.ct_final - self.ct_initial

    def get_flatten_index(self, idx):
        ''' maps the bumpy index of the node (index of each axis) to
            the flatten index of the node, i.e. the node number.

        Parameters
        ----------
        idx: tuple
            bumpy index of the node

        Returns
        -------
        int
            flatten index of the node
        '''
        assert type(idx) == tuple, ''
        assert len(idx) == self.sde.d, ''

        k = 0
        for i in range(self.sde.d):
            assert 0 <= idx[i] <= self.sde.Nx[i] - 1, ''
            Nx_prod = 1
            for j in range(i+1, self.sde.d):
                Nx_prod *= self.sde.Nx[j]
            k += idx[i] * Nx_prod

        return k

    def get_bumpy_index(self, k):
        ''' maps the flatten index of the node (node number) to
            the bumpy index of the node.

        Parameters
        ----------
        k: int
            flatten index of the node

        Returns
        -------
        tuple
            bumpy index of the node
        '''
        #assert type(k) == int, ''
        assert 0 <= k <= self.sde.Nh -1, ''

        idx = [None for i in range(self.sde.d)]
        for i in range(self.sde.d):
            Nx_prod = 1
            for j in range(i+1, self.sde.d):
                Nx_prod *= self.sde.Nx[j]
            idx[i] = k // Nx_prod
            k -= idx[i] * Nx_prod
        return tuple(idx)

    def get_x(self, k):
        ''' returns the x-coordinate of the node k

        Parameters
        ----------
        k: int
            flatten index of the node

        Returns
        -------
        float
            point in the domain
        '''
        assert k in np.arange(self.sde.Nh), ''

        return self.sde.domain_h.reshape(self.sde.Nh, self.sde.d)[k]

    def get_flatten_idx_from_axis_neighbours(self, idx, i):
        ''' get flatten idx of the neighbours with respect to the i-th coordinate

        Parameters
        ----------
        idx: tuple
            bumpy index of the node
        i: int
            index of the ith coordinate

        Returns
        -------
        tuple
            (k_left, k_right)
        '''

        # find flatten index of left neighbour wrt the i axis
        if idx[i] == 0:
            k_left = None
        else:
            left_idx = list(idx)
            left_idx[i] = idx[i] - 1
            k_left = self.get_flatten_index(tuple(left_idx))

        # find flatten index of right neighbour wrt the i axis
        if idx[i] == self.sde.Nx[i] - 1:
            k_right = None
        else:
            right_idx = list(idx)
            right_idx[i] = idx[i] + 1
            k_right = self.get_flatten_index(tuple(right_idx))

        return (k_left, k_right)

    def solve_bvp(self):
        ''' solves bvp by using finite difference
        '''

        # start timer
        self.start_timer()

        # discretized step
        h = self.h

        # discretize domain
        self.sde.discretize_domain_2d(h)

        # assemble linear system of equations: A \Psi = b.
        A = sparse.lil_matrix((self.sde.Nh, self.sde.Nh))
        b = np.zeros(self.sde.Nh)

        for k in np.arange(self.sde.Nh):

            # get discretized domain index
            idx = self.get_bumpy_index(k)

            # get point
            #x = self.get_x(idx)
            x = self.get_x(k)

            # assemble matrix A and vector b on S
            if k not in self.sde.ts_idx and k not in self.sde.boundary_idx:

                # drift and diffusion at x
                drift = self.sde.drift(x)
                sigma = self.sde.diffusion

                A[k, k] = - (sigma**2 * self.sde.d) / h**2 - self.sde.f(x)

                # x-axis neighbours
                k_left, k_right = self.get_flatten_idx_from_axis_neighbours(idx, i=0)
                A[k, k_left] = sigma**2 / (2 * h**2) - drift[0] / (2 * h)
                A[k, k_right] = sigma**2 / (2 * h**2) + drift[0] / (2 * h)

                # y-axis neighbours
                k_left, k_right = self.get_flatten_idx_from_axis_neighbours(idx, i=1)
                A[k, k_left] = sigma**2 / (2 * h**2) - drift[1] / (2 * h)
                A[k, k_right] = sigma**2 / (2 * h**2) + drift[1] / (2 * h)

            # impose condition on ∂S
            elif k in self.sde.ts_idx and not k in self.sde.boundary_idx:
                A[k, k] = 1
                b[k] = np.exp(- self.sde.g(x))

            # stability condition on the boundary: Psi should be flat
            elif k in self.sde.boundary_idx:
                neighbour_counter = 0

                if k in self.sde.boundary_x_idx:

                    # add neighbour
                    k_left, k_right = self.get_flatten_idx_from_axis_neighbours(idx, i=0)
                    if k_left is not None:
                        A[k, k_left] = - 1
                    elif k_right is not None:
                        A[k, k_right] = - 1

                    # update counter
                    neighbour_counter += 1

                if k in self.sde.boundary_y_idx:

                    # add neighbour
                    k_left, k_right = self.get_flatten_idx_from_axis_neighbours(idx, i=1)
                    if k_left is not None:
                        A[k, k_left] = - 1
                    elif k_right is not None:
                        A[k, k_right] = - 1

                    # update counter
                    neighbour_counter += 1

                # normalize
                A[k, k] = neighbour_counter

        # solve linear system and save
        psi = linalg.spsolve(A.tocsc(), b)
        self.psi = psi.reshape(self.sde.Nx)
        self.solved = True

        # stop timer
        self.stop_timer()

    def compute_value_function(self):
        ''' computes the value function
                phi = - log (psi)
        '''
        self.value_function =  - np.log(self.psi)

    def compute_optimal_control(self):
        ''' computes by finite differences the optimal control
                u_opt = - sigma ∇ value_f
        '''
        assert hasattr(self, 'value_function'), ''
        assert self.value_function.ndim == self.sde.d, ''

        # diffusion term
        sigma = self.sde.diffusion

        # preallocate optimal control
        self.u_opt = np.zeros(self.sde.Nx + (self.sde.d, ))

        for i in range(self.sde.d):

            # idx value f type
            value_f_k_plus_idx = [slice(self.sde.Nx[0]), slice(self.sde.Nx[1])]
            value_f_k_minus_idx = [slice(self.sde.Nx[0]), slice(self.sde.Nx[1])]

            # idx u type
            u_k_idx = [slice(self.sde.Nx[0]), slice(self.sde.Nx[1]), i]
            u_0_idx = [slice(self.sde.Nx[0]), slice(self.sde.Nx[1]), i]
            u_1_idx = [slice(self.sde.Nx[0]), slice(self.sde.Nx[1]), i]
            u_N_minus_idx = [slice(self.sde.Nx[0]), slice(self.sde.Nx[1]), i]
            u_N_idx = [slice(self.sde.Nx[0]), slice(self.sde.Nx[1]), i]

            value_f_k_plus_idx[i] = slice(2, self.sde.Nx[i])
            value_f_k_minus_idx[i] = slice(0, self.sde.Nx[i] - 2)
            u_k_idx[i] = slice(1, self.sde.Nx[i] - 1)
            u_0_idx[i] = 0
            u_1_idx[i] = 1
            u_N_minus_idx[i] = self.sde.Nx[i] - 2
            u_N_idx[i] = self.sde.Nx[i] - 1

            # generalized central difference
            self.u_opt[tuple(u_k_idx)] = - sigma *(
                self.value_function[tuple(value_f_k_plus_idx)]
              - self.value_function[tuple(value_f_k_minus_idx)]
            ) / (2 * self.sde.h)
            self.u_opt[tuple(u_0_idx)] = self.u_opt[tuple(u_1_idx)]
            self.u_opt[tuple(u_N_idx)] = self.u_opt[tuple(u_N_minus_idx)]

    def save(self):
        ''' saves some attributes as arrays into a .npz file
        '''
        # create data dictionary
        data = {
            'h': self.sde.h,
            'domain_h': self.sde.domain_h,
            'Nx': self.sde.Nx,
            'Nh': self.sde.Nh,
            'psi': self.psi,
            'value_function': self.value_function,
            'u_opt': self.u_opt,
            'ct': self.ct,
        }
        # save arrays in a npz file
        save_data(data, self.rel_dir_path)

    def load(self):
        ''' loads the saved arrays and sets them as attributes back
        '''
        data = load_data(self.rel_dir_path)
        try:
            for attr_name in data.keys():

                # get attribute from data
                attr = data[attr_name]

                # Controlled SDE attribute
                if attr_name in ['h', 'domain_h', 'Nx', 'Nh']:

                    # if attribute exists check if they are the same
                    if hasattr(self.sde, attr_name):
                        assert getattr(self.sde, attr_name) == attr

                    # if attribute does not exist save attribute
                    else:
                        setattr(self.sde, attr_name, attr)

                # hjb solver attribute
                else:

                    # if attribute exists check if they are the same
                    if hasattr(self, attr_name):
                        assert getattr(self, attr_name) == attr

                    # if attribute does not exist save attribute
                    else:
                        setattr(self, attr_name, attr)

            return True

        except:
            print('Attribute to load already exists and does not match')
            return False

    def get_psi_at_x(self, x):
        ''' evaluates solution of the BVP at x

        Parameters
        ----------
        x: array
            point in the domain

        Returns
        -------
        float
            psi at x
        '''
        # get index of x
        idx = self.sde.get_idx(x)

        # evaluate psi at x
        return self.psi[idx] if hasattr(self, 'psi') else None

    def get_value_function_at_x(self, x):
        ''' evaluates the value function at x

        Parameters
        ----------
        x: array
            point in the domain

        Returns
        -------
        float
            value function at x
        '''
        # get index of x
        idx = self.sde.get_idx(x)

        # evaluate value function at x
        return self.value_function[idx] if hasattr(self, 'value_function') else None

    def get_u_opt_at_x(self, x):
        ''' evaluates the optimal control at x

        Parameters
        ----------
        x: array
            point in the domain

        Returns
        -------
        array
            optimal control at x
        '''
        # get index of x
        idx = self.sde.get_idx(x)

        # evaluate optimal control at x
        return self.u_opt[idx] if hasattr(self, 'u_opt') else None

    def get_perturbed_potential_and_drift(self):
        ''' computes the potential, bias potential, controlled potential, gradient,
            controlled drift
        '''

        # flatten domain_h
        x = self.sde.domain_h.reshape(self.sde.Nh, self.sde.d)

        # diffusion term
        sigma = self.sde.diffusion

        # potential, bias potential and tilted potential
        self.V = self.sde.potential(x).reshape(self.sde.Nx)
        self.bias_potential = (sigma**2) * self.value_function
        self.perturbed_potential = self.V + self.bias_potential

        # gradient and tilted drift
        self.dV = self.sde.gradient(x).reshape(self.sde.domain_h.shape)
        self.perturbed_drift = - self.dV + sigma * self.u_opt

    def write_report(self, x):
        ''' writes the hjb solver parameters and the value of the solution at x

        Parameters
        ----------
        x: array
            point in the domain
        '''
        from sde_hjb_solver.utils import get_time_in_hms

        # space discretization
        print('\n space discretization')
        print('h = {:2.4f}'.format(self.sde.h))
        print('N_h = {:d}'.format(self.sde.Nh))

        # psi, value function and control
        print('\n psi, value function and optimal control at x')

        print('x: {:2.3f}'.format(x))
        psi = self.get_psi_at_x(x)
        value_f = self.get_value_function_at_x(x)
        u_opt = self.get_u_opt_at_x(x)

        if psi is not None:
            print('psi(x) = {:2.4e}'.format(psi))

        if value_f is not None:
            print('value_f(x) = {:2.4e}'.format(value_f))

        if u_opt is not None:
            print('u_opt(x): {:2.3f}'.format(u_opt))

        # maximum value of the control
        print('\n maximum value of the optimal control')

        idx_u_max = np.argmax(self.u_opt)
        x_u_max = self.get_x(idx_u_max)
        u_opt_max = self.u_opt[idx_u_max]
        print('argmax_x u_opt(x): {:2.3f}'.format(x_u_max))
        print('max_x u_opt(x): {:2.3f}'.format(u_opt_max))

        # computational time
        h, m, s = get_time_in_hms(self.ct)
        print('\nComputational time: {:d}:{:02d}:{:02.2f}\n'.format(h, m, s))

    def plot_2d_psi(self):
        fig, ax = plt.subplots()
        ax.set_title(r'Estimation of $\Psi(x)$')
        ax.set_xlabel(r'$x_1$')
        ax.set_ylabel(r'$x_2$')
        ax.set_xlim(self.sde.domain[0])
        ax.set_ylim(self.sde.domain[1])

        # contour f
        cs = ax.contourf(
            self.sde.domain_h[:, :, 0],
            self.sde.domain_h[:, :, 1],
            self.psi,
            extend='both',
        )

        # colorbar
        cbar = fig.colorbar(cs)

        plt.show()

    def plot_2d_value_function(self):
        fig, ax = plt.subplots()
        ax.set_title(r'Estimation of $\Phi(x)$')
        ax.set_xlabel(r'$x_1$')
        ax.set_ylabel(r'$x_2$')
        ax.set_xlim(self.sde.domain[0])
        ax.set_ylim(self.sde.domain[1])

        # contour f
        cs = ax.contourf(
            self.sde.domain_h[:, :, 0],
            self.sde.domain_h[:, :, 1],
            self.value_function,
            extend='both',
        )

        # colorbar
        cbar = fig.colorbar(cs)

        plt.show()

    def plot_2d_perturbed_potential(self):
        fig, ax = plt.subplots()
        ax.set_title(r'Perturbed potential $(V + V_{bias})(x)$')
        ax.set_xlabel(r'$x_1$')
        ax.set_ylabel(r'$x_2$')
        ax.set_xlim(self.sde.domain[0])
        ax.set_ylim(self.sde.domain[1])

        # contour f
        cs = ax.contourf(
            self.sde.domain_h[:, :, 0],
            self.sde.domain_h[:, :, 1],
            self.perturbed_potential,
            extend='both',
        )

        # colorbar
        cbar = fig.colorbar(cs)

        plt.show()

    def plot_2d_control(self, scale=None, width=0.005):
        from matplotlib import colors, cm

        fig, ax = plt.subplots()
        ax.set_title(r'Optimal control $u^*(x)$')
        ax.set_xlabel(r'$x_1$')
        ax.set_ylabel(r'$x_2$')
        ax.set_xlim(self.sde.domain[0])
        ax.set_ylim(self.sde.domain[1])

        X = self.sde.domain_h[:, :, 0]
        Y = self.sde.domain_h[:, :, 1]
        U = self.u_opt[:, :, 0]
        V = self.u_opt[:, :, 1]

        # set colormap
        colormap = cm.get_cmap('viridis_r', 100)
        colormap = colors.ListedColormap(
            colormap(np.linspace(0.20, 0.95, 75))
        )

        # initialize norm object and make rgba array
        C = np.sqrt(U**2 + V**2)
        norm = colors.Normalize(vmin=np.min(C), vmax=np.max(C))
        sm = cm.ScalarMappable(cmap=colormap, norm=norm)

        # quiver
        quiv = ax.quiver(
            X,
            Y,
            U,
            V,
            C,
            cmap=colormap,
            angles='xy',
            scale_units='xy',
            scale=scale,
            width=width,
        )

        # colorbar
        fig.colorbar(sm)

        plt.show()


# -*- coding: utf-8 -*-
"""
Created on Wed Jan 18 14:06:09 2023

@author: iant

LNO SPECTRAL CALIRATION PART 2 - READ IN FULLSCAN ANALYSIS RESULTS AND OUTPUT SPECTRAL AND TEMPERATURE COEFFICIENTS

spectral calibration is of the form
v / m = pxt^2 * A + pxt * B + C

pxt = px + (t * D) where px = the pixel number, t is the temperature and A, B, C and D are the coefficients to fit

First, find D: Take data for each order/line combination at a range of temperatures, the gradient is D

Then interpolate to find the zero celsius value i.e where pxt = px for each order/line combination
Then polyfit v/m versus the pxt pixels calculated
"""

import numpy as np
import matplotlib.pyplot as plt

filename = r"output\lno_spectral_calibration\LNO_fullscan_spectral_analysis3.txt" #3rd version i.e. after the 2nd iteration

input_dict = {"orders":[], "ts":[], "nus":[], "pxs":[], "nu_orders":[], "delta_nus":[]}

with open(filename, "r") as f:
    lines = f.readlines()

    for line in lines[1:]:
        line_split = line.split("\t")

        input_dict["orders"].append(float(line_split[1]))
        input_dict["ts"].append(float(line_split[2]))
        input_dict["nus"].append(float(line_split[3]))
        input_dict["pxs"].append(float(line_split[4]))
        input_dict["delta_nus"].append(float(line_split[5]))

        input_dict["nu_orders"].append(input_dict["nus"][-1]/input_dict["orders"][-1])

mean_delta_nu = np.mean(input_dict["delta_nus"])
print("mean delta nu = %0.3f" %mean_delta_nu)
std_delta_nu = np.std(input_dict["delta_nus"])
print("standard deviation delta nu = %0.3f" %std_delta_nu)

spectrum_px = np.arange(320)

unique_nu_orders = list(set(input_dict["nu_orders"]))

gradient_dict = {"orders":[], "nu_orders":[], "gradients":[], "px_zero_celsius":[]}

plt.figure(figsize=(12, 6), constrained_layout=True)
plt.grid()
plt.xlabel("Pixel number of absorption")
plt.ylabel("Instrument temperature C")
plt.title("LNO occultation fullscan: absorption shifts with temperature")
for unique_nu_order in unique_nu_orders:
    ixs = [i for i, v in enumerate(input_dict["nu_orders"]) if v == unique_nu_order]

    # print(len(ixs))
    if len(ixs) > 4:
        chosen_pxs = [input_dict["pxs"][i] for i in ixs]
        chosen_ts = [input_dict["ts"][i] for i in ixs]

        coeffs = np.polyfit(chosen_pxs, chosen_ts, 1)
        fit_ts = np.polyval(coeffs, chosen_pxs)

        "t = coeffs[0] * px + coeffs[1]"
        px_zero_celsius = -coeffs[1]/coeffs[0]

        # plt.scatter([i - np.mean(chosen_pxs) for i in chosen_pxs], chosen_ts)
        # plt.plot([i - np.mean(chosen_pxs) for i in chosen_pxs], fit_ts)
        plt.scatter(chosen_pxs, chosen_ts)
        plt.plot(chosen_pxs, fit_ts)

        gradient_dict["orders"].append(input_dict["orders"][ixs[0]])
        gradient_dict["nu_orders"].append(unique_nu_order)
        gradient_dict["gradients"].append(-1./coeffs[0])
        gradient_dict["px_zero_celsius"].append(px_zero_celsius)

plt.savefig("LNO_occultation_fullscan_absorption_shifts_with_temperature.png")

        # stop()
plt.figure(figsize=(12, 6), constrained_layout=True)
plt.grid()
plt.xlabel("Diffraction order")
plt.ylabel("D coefficient")
plt.title("LNO occultation fullscan: pixel shift per degree temperature change")
plt.scatter(gradient_dict["orders"], gradient_dict["gradients"])
gradient_coeffs = np.polyfit(gradient_dict["orders"], gradient_dict["gradients"], 1)
gradient_fits = np.polyval(gradient_coeffs, gradient_dict["orders"])
plt.plot(gradient_dict["orders"], gradient_fits)

plt.savefig("LNO_occultation_fullscan_pixel_shift_per_degree.png")

mean_gradient = np.mean(gradient_fits)

print("D = ", mean_gradient)

plt.figure(figsize=(12, 6), constrained_layout=True)
plt.grid()
plt.xlabel("Pixel number of absorption at 0 degrees C")
plt.ylabel("Absorption wavenumber / diffraction order")
plt.title("LNO occultation fullscan: pixel to wavenumber calibration")
plt.scatter(gradient_dict["px_zero_celsius"], gradient_dict["nu_orders"])
px_nu_coeffs = np.polyfit(gradient_dict["px_zero_celsius"], gradient_dict["nu_orders"], 2)
fit_px_nu_order = np.polyval(px_nu_coeffs, gradient_dict["px_zero_celsius"])
plt.plot(gradient_dict["px_zero_celsius"], fit_px_nu_order)

plt.savefig("LNO_occultation_fullscan_pixel_to_wavenumber_calibration.png")

print("A = ", px_nu_coeffs[0])
print("B = ", px_nu_coeffs[1])
print("C = ", px_nu_coeffs[2])

def it23_waven(diffraction_order, temperature):
    """LNO spectral calibration Ian January 22"""
    A =  3.942616548909985e-09
    B =  0.0005612789704190822
    C =  22.471896082567795
    D =  -0.851152813531944

    cfpixel = [A, B, C]

    p0_nu = D * temperature #px/°C * instrument temperature
    px_t_offset = np.arange(320.) + p0_nu
    nu  = np.polyval(cfpixel, px_t_offset) * diffraction_order

    return nu #wavenumbers of pixels

# A = 3.32e-8
# B = 5.480e-4
# C = 22.4701
# D = -0.8276

# for i in range(len(orders)):

#     nu_calc = it23_waven(order, inter_temp, px_in, A, B, C, D)

#     delta_nu = np.abs(nu_calc - nu)