import tcod
import random
import copy
import constants as const
import entity
import render
import numpy as np
import random_loot as rloot

class Room:
    """
    A room! Wow.
    """
    def __init__(self, x, y, w, h):
        self.x = x # upper left point
        self.y = y # upper left point
        self.w = w
        self.h = h
        self.n_loot = 0
        self.neighbors = []

class GameMap:
    def __init__(self, width, height, con, show_map=False):
        self.sample = None
        self.width = width
        self.con = con
        self.height = height
        self.tcod_empty_map = tcod.map.Map(self.width, self.height)
        self.show_map = show_map
        for x in range(self.width):
            for y in range(self.height):
                self.tcod_empty_map.transparent[y,x] = True
                self.tcod_empty_map.walkable[y,x] = True

    def get_sample(self):
        """
        Used to color the walls and the floor
        """
        ogrid = [np.arange(self.width, dtype=np.float32), np.arange(self.height, dtype=np.float32)]

        noise = tcod.noise.Noise(
            dimensions=2,
            algorithm=tcod.NOISE_PERLIN,
            implementation=tcod.noise.TURBULENCE,
            hurst=0.9,
            lacunarity=1.6,
            octaves=5)
        min_lum = 0.5
        max_lum = 1
        self.sample = noise.sample_ogrid(ogrid)*(max_lum-min_lum) + min_lum

    def add_loot(self, turns, player, entities):
        """
        Add loot to a new level
        """
        n_slot = {}
        for fslot in const.FeatureSlot:
            # 2d4 features of each slot
            n_slot[fslot] = sum([random.randint(1,4) for i in range(2)])
        for wslot in const.WeaponSlot:
            # 1d4 weapons of each slot
            n_slot[wslot] = sum([random.randint(1,4) for i in range(1)])

        for slot in n_slot:
            n_generated = n_slot.get(slot)
            for n in range(n_generated):
                n_try = 50
                while n_try > 0:
                    n_try -= 1
                    # loot is more probable in room with low arity
                    arity = random.choice([1,1,1,2,2,3,4,5])
                    rlist = [r for r in self.rooms_with_arity(arity) if r.n_loot < const.max_item_per_room]
                    if rlist:
                        room = random.choice(rlist)
                        (x,y) = self.random_cell_in_room(room)
                        if not self.tiles[x][y].item:
                            room.n_loot += 1
                            self.tiles[x][y].put_item(rloot.get_random_loot(slot, turns, player), entities)
                            break

    def rooms_with_arity(self, max_arity):
        """
        Return the list of room with at most max_arity neighbors
        """
        return [r for r in self.room_list if len(r.neighbors) <= max_arity]

    def make_boss_map(self, turns, entities, player):
        """
        An arena
        """
        self.get_sample()
        self.tiles = [[entity.Tile(x,y,color_coeff=self.sample[x][y]) for y in range(self.height)] for x in range(self.width)]
        self.tcod_map = tcod.map.Map(self.width, self.height)
        center_x = self.width / 2
        center_y = self.height / 2
        a = 30
        b = 15
        self.room_list = [Room(int(center_x - a), int(center_y - b), 2*a, 2*b)]
        for x in range(self.width):
            for y in range(self.height):
                if ((x - center_x)/a)**2 + ((y - center_y)/b)**2 <= 1: # ellipse equation
                    self.set_unblocked(x, y)
        (player.x, player.y) = (int(center_x / 2), int(center_y))
        boss = entity.Boss(int(center_x * 1.5), int(center_y))
        entities.append(boss)
        turns.add_turn(boss.speed_mov, const.TurnType.ENEMY, boss)
        self.recompute_fov(player.x, player.y)
        return boss

    def make_map_bsp(self, turns, entities, player):
        self.get_sample()
        self.tiles = [[entity.Tile(x,y,color_coeff=self.sample[x][y]) for y in range(self.height)] for x in range(self.width)]
        self.room_list = None
        self.tcod_map = tcod.map.Map(self.width, self.height)
        map_width = self.width
        map_height = self.height
        if self.show_map:
            for x in range(map_width):
                for y in range(map_height):
                    self.tiles[x][y].is_seen = True
        # we garantee a wall on the north and the west
        # this is necessary due to the generation the room
        bsp = tcod.bsp.BSP(1,1,map_width-1, map_height-1)
        bsp.split_recursive(6,6,6,1,1)
        self.room_list = self.recursive_make_rooms(bsp)

        # After the BSP generation, the dungeon is a tree
        # Create some loops
        rlist = self.rooms_with_arity(2)
        for i in range(6):
            for j in range(10):
                c = random.choice(range(len(rlist)))
                best = self.closest_rooms([rlist[c]], self.room_list)
                if best:
                    astar = tcod.path.AStar(self.tcod_map)
                    score_tuple = None
                    best_tuple = []
                    for tuple_param in best:
                        (x1, y1, x2, y2, _, _) = tuple_param
                        path = astar.get_path(x1, y1, x2, y2)
                        tmp_score = int(len(path)/3)
                        if not score_tuple or tmp_score > score_tuple:
                            score_tuple = tmp_score
                            best_tuple = [tuple_param]
                        elif tmp_score == score_tuple:
                            best_tuple.append(tuple_param)
                    self.connect_rooms(random.choice(best_tuple))
                    del rlist[c]
                    break

        # Initialization
        (player.x, player.y) = self.random_cell()
        (x, y) = self.random_cell()
        self.place_stairs(x,y)

        (x, y) = self.random_cell()
        self.place_boss_stairs(x,y)
        # self.place_boss_stairs(player.x,player.y) # DEBUG
        self.add_loot(turns, player, entities)
        self.recompute_fov(player.x, player.y)

    def recompute_fov(self, x, y, light_walls=True, radius=0):
        self.tcod_map.compute_fov(x, y, algorithm=2, radius=radius, light_walls=light_walls)

    def is_visible(self, x, y):
        return self.tcod_map.fov[y,x]

    def spawn_boss(self, entities, fslot, level, player):
        for i in range(50):
            (x,y) = self.random_cell()
            if not any([entity for entity in entities if entity.x == x and entity.y == y]):
                if (fslot == const.FeatureSlot.i and level >= 3) or (fslot != const.FeatureSlot.i and level >= 2):
                    class_name = fslot.value.get("bug_class")
                    the_class = getattr(entity, class_name)
                    monster = the_class(x, y, level, player.fequiped.get(fslot), fslot)
                else:
                    monster = entity.Monster(x, y, level, None, fslot)
                entities.append(monster)
                return monster
        return None

    def spawn(self, entities, feature):
        # We try at most 50 times to spawn it
        for i in range(50):
            (x,y) = self.random_cell()
            if not self.is_visible(x,y) and not any([entity for entity in entities if entity.x == x and entity.y == y]):
                level = random.randint(1, 3)
                if feature.n_bugs[level - 1] < const.n_bugs_max[feature.level - 1][level - 1]:
                    # mapgen bug are OP. Give their abitity to level 3 bug only
                    if (feature.fslot == const.FeatureSlot.i and level >= 3) or (feature.fslot != const.FeatureSlot.i and level >= 2):
                        class_name = feature.fslot.value.get("bug_class")
                        the_class = getattr(entity, class_name)
                        monster = the_class(x, y, level, feature)
                    else:
                        monster = entity.Monster(x, y, level, feature)
                    entities.append(monster)
                    return monster
        return None

    def iterator_perimeter_room(self, r):
        for x in range(r.x, r.x + r.w):
            yield (x, r.y)
            yield (x, r.y + r.h - 1)
        # y has a shorter range because the corners are already yielded
        for y in range(r.y + 1, r.y + r.h - 1):
            yield (r.x, y)
            yield (r.x + r.w - 1, y)

    def closest_rooms(self, l1, l2):
        best = []
        score_best = None
        for r1 in l1:
            for r2 in l2:
                if r1 != r2 and r1 not in r2.neighbors:
                    for (x1, y1) in self.iterator_perimeter_room(r1):
                        for (x2, y2) in self.iterator_perimeter_room(r2):
                            dx = abs(x1-x2)
                            dy = abs(y1-y2)
                            # This is not a hack. It is… hand-crafted mapgen
#                            if dx >= 4 and dy >= 4:
#                                score = max(abs(x1-x2),abs(y1-y2)) # Chebyshev distance
#                            else:
                            score = abs(x1-x2) + abs(y1-y2) # Manhattan distance
                            if score_best == None or score < score_best:
                                score_best = score
                                best = [(x1,y1,x2,y2,r1,r2)]
                            elif score == score_best:
                                best.append((x1,y1,x2,y2,r1,r2))
        return best

    def random_cell(self):
        return self.random_cell_in_room(random.choice(self.room_list))

    def random_cell_in_room(self, r):
        while True:
            x = random.randrange(r.x, r.x + r.w)
            y = random.randrange(r.y, r.y + r.h)
            if self.is_floor(x,y):
                return (x,y)

    def recursive_make_rooms(self, bsp):
        if not bsp.children:
            w = random.randrange(max(3,int(bsp.w/3)),bsp.w-2)
            h = random.randrange(max(3,int(bsp.h/3)),bsp.h-2)
            upper_left_x = random.randrange(bsp.x, bsp.x + bsp.w - w)
            upper_left_y = random.randrange(bsp.y, bsp.y + bsp.h - h)

            for x in range(0,w):
                for y in range(0,h):
                    self.set_unblocked(upper_left_x + x, upper_left_y + y)

            # Sometimes, add a central pillar
            if (w % 2) == 1 and (h % 2) == 1:
                if random.randrange(0,10) == 0:
                    center_x = upper_left_x + int((w-1)/2)
                    center_y = upper_left_y + int((h-1)/2)
                    self.set_blocked(center_x, center_y)

            # And rarely a big one (rare because big rooms aren't common)
            if (w % 2) == 0 and (h % 2) == 0 and w >= 10 and h >= 10 and random.randrange(0,2) == 0:
                center_x = upper_left_x + int(w/2) - 1
                center_y = upper_left_y + int(h/2) - 1
                for x in range(0,2):
                    for y in range(0,2):
                        self.set_blocked(center_x + x, center_y + y)

            return [Room(upper_left_x, upper_left_y, w, h)]

        else:
            l1 = self.recursive_make_rooms(bsp.children[0])
            l2 = self.recursive_make_rooms(bsp.children[1])
            # it is garanteed to connect
            self.connect_rooms(random.choice(self.closest_rooms(l1,l2)))
            return l1+l2

    def connect_rooms(self, tuple_param, force=False):
        (x1, y1, x2, y2, r1, r2) = tuple_param
        r1.neighbors.append(r2)
        r2.neighbors.append(r1)
        door_chance = 4
        if x1 == x2:
            if y1 > y2:
                y1 -= 1
                y2 += 1
            else:
                y1 += 1
                y2 -= 1
            self.create_v_tunnel(y1, y2, x1)
            if random.randint(0,door_chance) == 0:
                self.place_door(x1, y1)
            elif random.randint(0,door_chance) == 0:
                self.place_door(x2, y2)
        elif y1 == y2:
            if x1 > x2:
                x1 -= 1
                x2 += 1
            else:
                x1 += 1
                x2 -= 1
            self.create_h_tunnel(x1, x2, y1)
            if random.randint(0,door_chance) == 0:
                self.place_door(x1, y1)
            elif random.randint(0,door_chance) == 0:
                self.place_door(x2, y2)
#        elif abs(x1-x2) < 3 or abs(y1-y2) < 3:
        else:
            if random.randint(0, 1) == 1:
                if x1 > x2:
                    x1 -= 1
                else:
                    x1 += 1
                if y1 > y2:
                    y2 += 1
                    y3 = y1 - 1
                else:
                    y2 -= 1
                    y3 = y1 + 1

                self.create_h_tunnel(x1, x2, y1)
                self.create_v_tunnel(y3, y2, x2)
                if random.randint(0,door_chance) == 0 and abs(x1-x2) > 1:
                    self.place_door(x1, y1)
                elif random.randint(0,door_chance) == 0 and abs(y1-y2) > 1:
                    self.place_door(x2, y2)

            else:
                if x1 > x2:
                    x2 += 1
                    x3 = x1 - 1
                else:
                    x2 -= 1
                    x3 = x1 + 1
                if y1 > y2:
                    y1 -= 1
                else:
                    y1 += 1

                self.create_v_tunnel(y1, y2, x1)
                self.create_h_tunnel(x3, x2, y2)
                if random.randint(0,door_chance) == 0 and abs(y1-y2) > 1:
                    self.place_door(x1, y1)
                elif random.randint(0,door_chance) == 0 and abs(x1-x2) > 1:
                    self.place_door(x2, y2)

    def create_h_tunnel(self, x1, x2, y):
        for x in range(min(x1, x2), max(x1, x2) + 1):
            self.set_unblocked(x,y)

    def create_v_tunnel(self, y1, y2, x):
        for y in range(min(y1, y2), max(y1, y2) + 1):
            self.set_unblocked(x,y)

    def get_copy_map(self):
        return copy.deepcopy(self.tcod_map)

    def get_copy_empty_map(self):
        return copy.deepcopy(self.tcod_empty_map)

    def set_tile_type(self, x, y, ttype):
        self.tiles[x][y] = entity.Tile(x, y, color_coeff=self.sample[x][y], ttype=ttype)
        if self.show_map:
            self.tiles[x][y].is_seen = True
        self.tcod_map.transparent[y,x] = ttype.value.get("transparent")
        self.tcod_map.walkable[y,x] = not ttype.value.get("collision")

    def is_over_map(self, x, y):
        return x >= 0 and y >= 0 and x < self.width and y < self.height

    def set_blocked(self, x, y):
        self.set_tile_type(x, y, const.TileType.WALL)

    def set_unblocked(self, x, y):
        self.set_tile_type(x, y, const.TileType.FLOOR)

    def place_door(self, x, y):
        self.set_tile_type(x, y, const.TileType.DOOR)

    def place_stairs(self, x, y):
        self.set_tile_type(x, y, const.TileType.STAIRS)

    def place_boss_stairs(self, x, y):
        self.set_tile_type(x, y, const.TileType.BOSS_STAIRS)

    def is_floor(self, x, y):
        return self.tiles[x][y].ftype == const.TileType.FLOOR

    def is_door(self, x, y):
        return self.tiles[x][y].ftype == const.TileType.DOOR

    def is_stairs(self, x, y):
        return self.tiles[x][y].ftype == const.TileType.STAIRS

    def is_boss_stairs(self, x, y):
        return self.tiles[x][y].ftype == const.TileType.BOSS_STAIRS

    def is_blocked(self, x, y):
        return not self.tcod_map.walkable[y,x]

    def drop_item_on_floor(self, player, entities, item, drop_key):
        if not self.tiles[player.x][player.y].item:
            player.remove_from_inventory(item, drop_key)
            return self.tiles[player.x][player.y].put_item(item, entities)

    def is_weapon_on_floor_directly_equipable(self, player):
        item = self.tiles[player.x][player.y].item
        if item and isinstance(item, entity.Weapon) and not player.wequiped.get(item.wslot):
            return True
        return False

    def get_item_on_floor(self, player, entities):
        if self.tiles[player.x][player.y].item:
            item = self.tiles[player.x][player.y].take_item(entities)
            key = player.add_to_inventory(item)
            return (item,key)

    def description_item_on_floor(self, player):
        """
        Get the name of the item on the floor where the player is
        """
        if self.tiles[player.x][player.y].item:
            return self.tiles[player.x][player.y].item.name
        return None

    def is_there_item_on_floor(self, player):
        """
        Is there an item on the floor, where the player is?
        """
        return self.tiles[player.x][player.y].item != None


import argparse
from dateutil import tz
from datetime import datetime

from spotipy import Spotify
import spotipy.util

from models import Play, Track, Album, Artist, PostgreSQLConnection

import settings

def set_timezone_to_datetime(datetime_to_set, timezone):
    return datetime_to_set.replace(tzinfo=tz.gettz(timezone))

def convert_played_at_from_response_to_datetime(played_at):
    try:
        return datetime.strptime(played_at, '%Y-%m-%dT%H:%M:%S.%fZ')
    except:
        # For the single moment where the played at time hits a full second
        return datetime.strptime(played_at, '%Y-%m-%dT%H:%M:%SZ')

def convert_datetime_from_timezone_to_timezone(datetime_to_convert, from_tz_code, to_tz_code):
    from_tz = tz.gettz(from_tz_code)
    to_tz = tz.gettz(to_tz_code)

    datetime_to_convert = datetime_to_convert.replace(tzinfo=from_tz)
    converted_datetime = datetime_to_convert.astimezone(to_tz)
    return converted_datetime

class SpotifyConnection(object):

    def __init__(self, user_data):
        self.user_name = user_data['user_name']
        token = spotipy.util.prompt_for_user_token(self.user_name,
                                                   scope='user-read-recently-played',
                                                   client_id=user_data['client_id'],
                                                   client_secret=user_data['client_secret'],
                                                   redirect_uri=user_data['redirect_uri'])
        self.client = Spotify(auth=token)
        self.db = self.init_db()

    def init_db(self):
        return PostgreSQLConnection()

    def get_artist(self, artist_id):
        artist = self.db.session.query(Artist).get(artist_id)
        if artist:
            return artist
        else:
            artist_response = self.client.artist(artist_id)
            artist = Artist()
            artist.artist_id = artist_id
            artist.artist_data = artist_response
            self.db.save_instance(artist)
            print("> Artist {} was not in database.".format(artist.artist_data['name']))
            return self.db.session.query(Artist).get(artist_id)

    def get_album(self, album_id):
        album = self.db.session.query(Album).get(album_id)
        if album:
            return album
        else:
            album_response = self.client.album(album_id)
            album = Album()
            album.album_data = album_response
            album.album_id = album_response['id']
            # Artists
            for album_artist_response in album_response['artists']:
                album.artists.append(self.get_artist(album_artist_response['id']))
            self.db.save_instance(album)
            print("> Album {} was not in database.".format(album.album_data['name']))
            return self.db.session.query(Album).get(album_id)

    def get_track(self, track_id):
        track = self.db.session.query(Track).get(track_id)
        if track:
            return track
        else:
            response = self.client.track(track_id)

            track = Track()
            track.track_id = track_id
            track.track_data = response
            # Album
            track.album = self.get_album(response['album']['id'])
            # Artists
            for artist_response in response['artists']:
                track.artists.append(self.get_artist(artist_response['id']))
            # Audio feature
            audio_feature_response = self.client.audio_features(track_id)[0]
            if audio_feature_response:  # Some tracks do not have audio features
                track.audio_feature_data = audio_feature_response
            print("> Track {} was not in database.".format(track.track_data['name']))
            self.db.save_instance(track)
            return self.db.session.query(Track).get(track_id)

    def get_play_from_played_at_utc_and_track_id(self, played_at_utc, track_id):
        played_at_utc = convert_played_at_from_response_to_datetime(played_at_utc)
        played_at_utc = set_timezone_to_datetime(played_at_utc, timezone='UTC')
        played_at_cet = convert_datetime_from_timezone_to_timezone(played_at_utc,
                                                                   from_tz_code='UTC',
                                                                   to_tz_code='CET')
        # Play
        play = Play()
        play.user_name = self.user_name
        play.played_at_utc_timestamp = played_at_utc.timestamp() * 1000
        play.played_at_utc = played_at_utc
        play.played_at_cet = played_at_cet
        play.day = played_at_cet.day
        play.month = played_at_cet.month
        play.year = played_at_cet.year
        play.hour = played_at_cet.hour
        play.minute = played_at_cet.minute
        play.second = played_at_cet.second
        play.day_of_week = played_at_cet.weekday()
        play.week_of_year = played_at_cet.date().isocalendar()[1]
        # Track
        track = self.get_track(track_id)
        play.track = track
        play.track_id = track_id
        return play

    def _get_play_tuples_from_response(self, response):
        plays = []
        for item in response['items']:
            play_tuple = (item['played_at'], item['track']['id'])
            plays.append(play_tuple)
        return plays

    def _get_play_tuples(self, limit=50, after=None):
        play_tuples = []
        response = self.client._get('me/player/recently-played', after=after, limit=limit)
        play_tuples.extend(self._get_play_tuples_from_response(response))

        while response and 'next' in response:
            response = self.client.next(response)
            if response:
                play_tuples.extend(self._get_play_tuples_from_response(response))

        return play_tuples

    def extract_plays(self):
        print("* Extracting latest plays of {}.".format(self.user_name))
        play_tuples = self._get_play_tuples()

        for played_at, track_id in play_tuples:
            play = self.get_play_from_played_at_utc_and_track_id(played_at, track_id)
            self.db.save_play(play)

class HoergewohnheitenManager(object):

    def __init__(self, spotify_user_data):
        self.spotify = SpotifyConnection(user_data=spotify_user_data)

    def process_hoergewohnheiten(self):
        self.spotify.extract_plays()

def process_hoergewohnheiten(user_name):
    print("***", user_name, "***")
    user_data = settings.SPOTIFY_USERS[user_name]
    mgr = HoergewohnheitenManager(user_data)
    mgr.process_hoergewohnheiten()

if __name__ == '__main__':
    print('''
 _     ___   ____  ___   __    ____  _       ___   _     _      _     ____  _  _____  ____  _
| |_| / / \ | |_  | |_) / /`_ | |_  \ \    // / \ | |_| | |\ | | |_| | |_  | |  | |  | |_  | |\ |
|_| | \_\_/ |_|__ |_| \ \_\_/ |_|__  \_\/\/ \_\_/ |_| | |_| \| |_| | |_|__ |_|  |_|  |_|__ |_| \|
''')

    print("Started at {}.".format(datetime.now()))

    # Argparse
    parser = argparse.ArgumentParser(description='Hoergewohnheiten')
    parser.add_argument('-u', dest='user_name')
    args = parser.parse_args()

    if args.user_name:
        process_hoergewohnheiten(args.user_name)
    else:
        for user_name in settings.SPOTIFY_USERS:
            process_hoergewohnheiten(user_name)

    print("Finished at {}.".format(datetime.now()))


"""
4.	Найти сумму n элементов следующего ряда чисел: 1 -0.5 0.25 -0.125 ...
Количество элементов (n) вводится с клавиатуры.
ЧЕРЕЗ ЦИКЛ
"""

num_el = int(input('enter 3 '))
summ = 0
range_num = 1
for i in range(num_el):
    summ += range_num
    range_num /= -2

print(summ)


import numpy as np
import torch

from affogato.affinities import compute_affinities
from torchvision.utils import make_grid
from inferno.extensions.criteria import SorensenDiceLoss

class ConcatDataset(torch.utils.data.Dataset):
    def __init__(self, *datasets):
        self.datasets = datasets
        self.lens = [len(ds) for ds in self.datasets]

        self.start_idx = np.cumsum(self.lens)
        self.start_idx[-1] = 0
        self.start_idx = np.roll(self.start_idx, 1)

    def __len__(self):
        return sum(self.lens)

    def __getitem__(self, index):
        ds_index = np.where(index - self.start_idx >= 0)[0][-1]
        item_index = index - self.start_idx[ds_index]
        return self.datasets[ds_index][item_index]

class DefaultDataset(torch.utils.data.Dataset):
    """ Simple default dataset for generating affinities
    from segmentation and mask.
    """
    patch_shape = [512, 512]  # TODO expose this and other parameters

    def to_affinities(self, seg, mask):
        seg[~mask] = 0
        affs, aff_mask = compute_affinities(seg, self.offsets, have_ignore_label=True)
        aff_mask = aff_mask.astype('bool')
        affs = 1. - affs

        mask_transition, aff_mask2 = compute_affinities(mask, self.offsets)
        mask_transition[~aff_mask2.astype('bool')] = 1
        aff_mask[~mask_transition.astype('bool')] = True
        return affs, aff_mask

    @staticmethod
    def estimate_n_samples(shape, patch_shape):
        # we estimate the number of samples by tiling shape with patch_shape
        crops_per_dim = [sh / float(cs) for sh, cs in zip(shape, patch_shape)]
        return int(np.prod(crops_per_dim))

    def __init__(self, raw, seg, mask_ids, offsets, transforms=None):
        self.raw = raw
        self.seg = seg
        self.mask_ids = mask_ids
        self.offsets = offsets
        self.transforms = transforms
        self.n_samples = self.estimate_n_samples(self.raw.shape, self.patch_shape)

    def __getitem__(self, index):

        # TODO sample so that we are biased towards the mask
        def sample_raw_seg_mask():
            offset = [np.random.randint(0, sh - csh) if sh > csh else 0
                      for sh, csh in zip(self.raw.shape, self.patch_shape)]
            bb = tuple(slice(off, off + csh) for off, csh in zip(offset, self.patch_shape))

            raw = self.raw[bb]
            seg = self.seg[bb]

            if self.transforms is not None:
                raw, seg = self.transforms(raw, seg)

            raw, seg = raw.copy(), seg.copy()
            mask = np.isin(seg, self.mask_ids)

            return raw, seg, mask

        raw, seg, mask = sample_raw_seg_mask()

        # TODO ensure that we have some in-mask area
        # # some arbitrary but very small pixel threshold
        # while mask.sum() < 25:
        #     raw, seg, mask = sample_raw_seg_mask()

        # add channel dim
        raw = raw[None]

        # make affs and aff_mask
        affs, aff_mask = self.to_affinities(seg, mask)

        return raw, affs, aff_mask

    def __len__(self):
        return self.n_samples

class MaskedLoss(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.criterion = SorensenDiceLoss()

    def forward(self, pred, y, mask):
        mask.requires_grad = False
        masked_prediction = pred * mask
        loss = self.criterion(masked_prediction, y)
        return loss

def default_training(proc_id, net, ds,
                     pipe, device, step):
    loader = torch.utils.data.DataLoader(ds, batch_size=1, num_workers=2)

    p_out, p_in = pipe

    optimizer = torch.optim.Adam(net.parameters(), lr=1e-4)
    loss = MaskedLoss()
    loss = loss.to(device)

    logger = torch.utils.tensorboard.SummaryWriter('./runs/imws')

    add_gradients = True
    log_frequency = 10

    net.train()
    while True:
        if p_out.poll():
            if not p_out.recv():
                p_in.send(step)
                break

        for x, y, mask in loader:
            x = x.to(device)
            y, mask = y.to(device), mask.to(device)

            optimizer.zero_grad()

            pred = net(x)
            pred.retain_grad()

            loss_val = loss(pred, y, mask)
            loss_val.backward()
            optimizer.step()

            logger.add_scalar("loss", loss_val.item(), step)

            step += 1
            if step % log_frequency == 0:
                print("Background training process iteration", step)
                x = x[0].detach().cpu()
                logger.add_image('input', x, step)
                y = y[0].detach().cpu()

                if add_gradients:
                    grads = pred.grad[0].detach().cpu()
                    grads -= grads.min()
                    grads /= grads.max()

                pred = torch.clamp(pred[0].detach().cpu(), 0.001, 0.999)
                tandp = [target.unsqueeze(0) for target in y]
                nrow = len(tandp)
                tandp.extend([p.unsqueeze(0) for p in pred])

                if add_gradients:
                    tandp.extend([grad.unsqueeze(0) for grad in grads])
                tandp = make_grid(tandp, nrow=nrow)
                logger.add_image('target_and_prediction', tandp, step)

                # for debugging
                # return x, y, pred, grads


import tensorflow as tf
import numpy as np
import time
#
from dataset.train import *
import convert
import neural
import test

class NeuralNetwork:
    def __init__(self):
        test_success = 0
        start = time.time()
        target = 0.01

        print('desfrdgthygju', len(data_input))

        [to_train, to_test, res_to_train, res_to_test] = convert.dataset_matrix_to_lists(data_input, data_output, 75, 30)

        for x in range(0, 1):
            [epoch, err] = neural.train(to_train, res_to_train)
            print('epoch:', epoch, 'mse:', err)
            if err < target:
                test_success = test_success + 1
            test.test_neural(epoch, to_test, res_to_test);

        print('time', time.time() - start)
        print('success ', test_success)

    def __del__(self):
        pass


"""
Given an integer array nums, find the subarray
with the largest sum, and return its sum.

Example 1:

Input: nums = [-2,1,-3,4,-1,2,1,-5,4]
Output: 6
Explanation: The subarray [4,-1,2,1] has the largest sum 6.

"""

"""
Time Complexity:- O(n)
Space Complexity:- O(1)
"""

class Solution:
    def maxSubArray(self, nums):
        # Initialize variables to keep track of the maximum sum and the current sum.
        max_sum = float("-inf")
        current_sum = 0

        # Iterate through the array to find the maximum subarray sum.
        for num in nums:
            # If the current subarray sum is negative, start a new subarray.
            # Otherwise, continue adding elements to the current subarray.
            current_sum = max(num, current_sum + num)

            # Update the maximum subarray sum if needed.
            max_sum = max(max_sum, current_sum)

        # Return the maximum subarray sum found.
        return max_sum


import inspect

from .namespace import Namespace

_matches_cache = {}

def matches(caller_parameters, callee_parameters):
    cache_key = ';'.join((caller_parameters, callee_parameters))  # pragma: no mutate
    cached_value = _matches_cache.get(cache_key, None)  # pragma: no mutate (mutation changes this to cached_value = None, which just slows down the code)
    if cached_value is not None:
        return cached_value

    caller = set(caller_parameters.split(',')) if caller_parameters else set()

    a, b, c = callee_parameters.split('|')
    required = set(a.split(',')) if a else set()
    optional = set(b.split(',')) if b else set()
    wildcard = (c == '*')

    if not required and not optional and wildcard:
        return False  # Special case to not match no-specification function "lambda **whatever: ..."

    if wildcard:
        result = caller >= required
    else:
        result = required <= caller <= required.union(optional)

    _matches_cache[cache_key] = result  # pragma: no mutate (mutation changes result to None which just makes things slower)
    return result

def get_callable_description(c):
    if getattr(c, '__name__', None) == '<lambda>':
        import inspect
        try:
            return 'lambda found at: `{}`'.format(inspect.getsource(c).strip())
        except OSError:
            pass
    return f'`{c}`'

def evaluate(func_or_value, __signature=None, __strict=False, **kwargs):
    if callable(func_or_value):
        if __signature is None:
            __signature = signature_from_kwargs(kwargs)

        callee_parameters = get_signature(func_or_value)
        if callee_parameters is not None and matches(__signature, callee_parameters):
            return func_or_value(**kwargs)
    if __strict and callable(func_or_value):
        assert (
            isinstance(func_or_value, Namespace)
            and 'call_target' not in func_or_value
        ), "Evaluating {} didn't resolve it into a value but strict mode was active, " \
           "the signature doesn't match the given parameters. " \
           "Note that you must match at least one keyword argument. " \
           "We had these arguments: {}".format(
            get_callable_description(func_or_value),
            ', '.join(kwargs.keys()),
        )
    return func_or_value

def evaluate_strict(func_or_value, __signature=None, **kwargs):
    # noinspection PyArgumentEqualDefault
    return evaluate(func_or_value, __signature=None, __strict=True, **kwargs)

def evaluate_recursive(func_or_value, __signature=None, __strict=False, **kwargs):
    if __signature is None:
        __signature = signature_from_kwargs(kwargs)  # pragma: no mutate

    if isinstance(func_or_value, dict):
        # The type(item)(** stuff is to preserve the original type
        return type(func_or_value)(**{k: evaluate_recursive(v, __signature=__signature, __strict=__strict, **kwargs) for k, v in dict.items(func_or_value)})

    if isinstance(func_or_value, list):
        return [evaluate_recursive(v, __signature=__signature, __strict=__strict, **kwargs) for v in func_or_value]

    if isinstance(func_or_value, set):
        return {evaluate_recursive(v, __signature=__signature, __strict=__strict, **kwargs) for v in func_or_value}

    return evaluate(func_or_value, __signature=__signature, __strict=__strict, **kwargs)

def evaluate_recursive_strict(func_or_value, __signature=None, **kwargs):
    """
    Like `evaluate_recursive` but won't allow un-evaluated callables to slip through.
    """
    # noinspection PyArgumentEqualDefault
    return evaluate_recursive(func_or_value, __signature=None, __strict=True, **kwargs)

def get_signature(func):
    """
        :type func: Callable
        :rtype: str
    """
    try:
        return object.__getattribute__(func, '__tri_declarative_signature')
    except AttributeError:
        pass

    try:
        names, _, varkw, defaults, _, _, _ = inspect.getfullargspec(func)
    except TypeError:
        return None

    first_arg_index = 1 if inspect.ismethod(func) else 0  # Skip self argument on methods

    number_of_defaults = len(defaults) if defaults else 0
    if number_of_defaults > 0:
        required = ','.join(sorted(names[first_arg_index:-number_of_defaults]))
        optional = ','.join(sorted(names[-number_of_defaults:]))
    else:
        required = ','.join(sorted(names[first_arg_index:]))
        optional = ''
    wildcard = '*' if varkw is not None else ''

    signature = '|'.join((required, optional, wildcard))
    try:
        object.__setattr__(func, '__tri_declarative_signature', signature)
    except TypeError:
        # For classes
        type.__setattr__(func, '__tri_declarative_signature', signature)
    except AttributeError:
        pass
    return signature

def signature_from_kwargs(kwargs):
    return ','.join(sorted(kwargs.keys()))


# -*- coding: utf-8 -*-
"""
Created on Tue Jun  6 16:06:45 2017

@author: Francesco
"""
import threading
import sys

import serial
import numpy as np
import time

import matplotlib.pyplot as plt

global PORT
global BAUD
global NUM_CHANNELS
global END_BUNDLE_BYTE
global BYTE_PER_CHANNEL
global BUNDLE_LENGTH

#BUNDLE SHAPE: |!|!|!|CH0_msb|CH0_lsb|ch1_msb|ch1_lsb|......|ch7_lsb|!|!|!|

PORT = "COM3"
BAUD = 115200
NUM_CHANNELS = 8
END_BUNDLE_BYTE = 3
BYTE_PER_CHANNEL = 2 #two bytes to represent int
BUNDLE_LENGTH = NUM_CHANNELS*BYTE_PER_CHANNEL

global ROWS
global COLS
ROWS = 4
COLS = 2

global ACQUISITION_TIME
ACQUISITION_TIME = 60
global SAMPLING_INTERVAL
SAMPLING_INTERVAL = 0.5
global n_MOVEMENT_TO_UNDERSTAND
n_MOVEMENT_TO_UNDERSTAND = 4 #up down left right
global movements
movements = ["up\n","down\n","left\n","right\n"]

class SerialReader(threading.Thread):

    def __init__(self, name, port_name, baud, data, event_run):
        threading.Thread.__init__(self)
        self.name = name
        self.port_name = port_name
        self.baud = baud
        self.event_run = event_run
        self.data = data
        print("Attempting to open port %s at baud %d" %(self.port_name,self.baud))
        self.port = serial.Serial(self.port_name,self.baud,timeout=1)
        if(self.port.isOpen()): print("Port Open")

    def run(self):

        start_time = time.time()
        running = True
        while(running):
            try:
                #actual decoding
                if(self.port.read(END_BUNDLE_BYTE).decode("raw_unicode_escape") == '!!!'):
                    temp = self.port.read(BUNDLE_LENGTH)
                    #print(temp)
                    for channel in range(0,NUM_CHANNELS):
                        self.data[channel] = (temp[channel*BYTE_PER_CHANNEL]<<8)|(temp[channel*BYTE_PER_CHANNEL + 1 ])
                    #allow the plotting thread to access the data
                    self.event_run.set()
                    self.event_run.clear()

                    time.sleep(SAMPLING_INTERVAL)

                total_elapsed = time.time() - start_time

                if(total_elapsed > ACQUISITION_TIME): #more than 30 seconds of acquisition
                    print("From %s, closing port"%self.name)
                    self.port.close()
                    running = False
            except KeyboardInterrupt:
                self.port.close()
                break

class DynamicPlotter(threading.Thread):

    def __init__(self,name,data,event_run):
        threading.Thread.__init__(self)
        # Scrive un file.
        self.out_file = open("test.txt","w")
        self.data = data
        self.event_run = event_run
        self.name = name
        self.number_of_acq_total = ACQUISITION_TIME/SAMPLING_INTERVAL
        self.number_of_acq_per_movement = self.number_of_acq_total/n_MOVEMENT_TO_UNDERSTAND

    def run(self):

        running = True
        counter_total = 0
        counter = 0
        while(running):
            self.event_run.wait()
            #print("From %s, writing to the file!"%self.name)
            self.out_file.write(str(self.data[0])) #only to avoid printing a coma
            for value in self.data[1:]:
                self.out_file.write(',')
                self.out_file.write(str(value))

            self.out_file.write('\n')

            if(counter == 0):
                print("Counter total:%d"%counter_total)
                index = int(counter_total/self.number_of_acq_per_movement)
                message = "Movement: %s"%movements[index]
                #why is this working?
                #let's suppose: counter=0 and counter_total =
                print("%s: %s"%(self.name,message))
                self.out_file.write(message)

            counter_total += 1
            counter += 1

            if(counter == self.number_of_acq_per_movement):
                #reset counter
                counter = 0

            #print("From %s, checking if set: %d!"%(self.name,self.event_close.is_set()))
            if(counter_total == self.number_of_acq_total ): #6 acquisitions,

                print("From %s, closing the file!"%self.name)
                self.out_file.close()
                running = False

if __name__ == "__main__":

    #matrix that holds values read by
    data = np.zeros(NUM_CHANNELS)

    event_run = threading.Event()

    try:

        s = SerialReader("Serial Reader",PORT,BAUD,data,event_run)
        p = DynamicPlotter("File maker",data,event_run)

        s.start()
        p.start()

        s.join()
        p.join()

    except KeyboardInterrupt:
        raise ValueError('Catch command to stop from keyboard')
        sys.exit(0)


import numpy as np
import subprocess
import pandas as pd
import os
import sys

def prompt_savefig(figure, filename):
    """
    Take a filename for a figure, and prompt the user to overwrite an already existing
    file with that name if needed
    """
    if os.path.exists(filename):
        prompt = input(f"{filename} already exists, overwrite? [y/N]> ")
        if not (prompt.lower() == "y" or prompt.lower() == "yes"):
            prompt_savefig(figure, input("Enter new filename: "))

    figure.savefig(filename)

def get_sorted_id_dirs(foldername):
    """
    Take a path to a multicore run folder and return a sorted list of id directories
    """

    return sorted(
            [directory for directory in os.listdir(foldername) if os.path.isdir(
                os.path.join(foldername, directory)
                )
                and directory[:2] == "id"
                and directory[2:].isdigit()],
            key=lambda directory: int(directory[2:]),
            )

def get_last_tab_file(id_foldername, id_directory_num=0, dim3=False):
    """
    Take a path to a(n) (id) directory, return the name of the tab file at the last
    time index
    """

    # final index on the tab files
    last_tab_time = get_last_tab_time(id_foldername)

    # build the tabfile name; the ones in the 0th id directory
    # aren't labeled by id, but the rest are
    id_filename_modifier = ""
    if id_directory_num > 0:
        id_filename_modifier = f"-id{id_directory_num}"

    dimstr = None
    if dim3:
        dimstr = "3d"
    else:
        dimstr = "2d"

    return f"Par_Strat{dimstr}{id_filename_modifier}.{last_tab_time:04}.tab"

def get_tab_file(idfolder, timestep_id):
    """
    Take a path to a(n) (id) directory, return the name of the tab file
    with the given timestep id
    """

    filelist = [f for f in os.listdir(idfolder) if os.path.isfile(os.path.join(idfolder, f))]
    tabfilelist = [f for f in filelist if "tab" in f]
    return [f for f in tabfilelist if f"{timestep_id:04}" in f][0]

def get_iddir_2d_chunk(id_dirs, id_offset, chunk2dsize):
    return [id_dir for id_dir in id_dirs
            if id_offset <= get_idnum_from_foldername(id_dir) < id_offset+chunk2dsize]

def get_idnum_from_foldername(id_foldername):
    return int(id_foldername[2:])

def get_tab_df_multicore(foldername, x1_blocks, x2_blocks, timestep=None, x3_block_offset=None):
    """
    Take a path to a multicore run folder and return a dataframe containing
    all of the tabfile data
    x1_blocks: Number of blocks (for MPI) in the x1 direction
    x2_blocks: Number of blocks (for MPI) in the x2 direction
    timestep: Tab file index
    """

    # This is a sorted list of the id directories in the parent folder
    id_directories = get_sorted_id_dirs(foldername)
    if x3_block_offset is not None:
        chunk2dsize = x1_blocks * x2_blocks
        id_offset = x3_block_offset * chunk2dsize
        id_directories = get_iddir_2d_chunk(id_directories, id_offset, chunk2dsize)

    i_offset = 0
    j_offset = 0
    previous_x2_block_index = 0
    master_df = None
    for id_directory in id_directories:
        id_directory_num = get_idnum_from_foldername(id_directory)
        tab_filename = None
        if timestep is None:
            tab_filename = get_last_tab_file(
                    os.path.join(foldername, id_directory),
                    id_directory_num,
                    dim3=(x3_block_offset is not None),
                    )
        else:
            tab_filename = get_tab_file(os.path.join(foldername, id_directory), timestep)

        df = get_df(os.path.join(foldername, id_directory, tab_filename))

        """
        Need to add an offset to the cell indices to put them in the right place, since
        cell indices are computed by athena relative to the CPU space (blocks), but we need
        them relative to the entire domain.
        """
        x1_block_index = id_directory_num % x1_blocks
        x2_block_index = np.floor(float(id_directory_num % chunk2dsize) / x1_blocks)

        # compute j_offset for this block
        if x2_block_index > previous_x2_block_index:
            j_offset += max(df["j-zone"]) - min(df["j-zone"]) + 1
        # add accumulated offset
        df["i-zone"] += i_offset
        df["j-zone"] += j_offset

        # compute k_offset for this block (if needed)
        if x3_block_offset is not None:
            df["k-zone"] += ( max(df["k-zone"]) - min(df["k-zone"]) + 1 ) * x3_block_offset

        """
        Compute the row (x2_block_index) and column (x1_block_index) of the current file's
        block. Athena creates id folders first along x1, then along x2, which is why these
        formulae work.
        """

        # compute i offset for next block
        if x1_block_index < x1_blocks - 1:
            i_offset += max(df["i-zone"]) - min(df["i-zone"]) + 1
        else:
            i_offset = 0

        previous_x2_block_index = x2_block_index

        if master_df is None:
            master_df = df
        else:
            master_df = master_df.append(df)

    return master_df

def get_tab_df_multicore_3d(foldername, x1_blocks, x2_blocks, x3_blocks, timestep=None):
    master_df = None
    for x3_index in range(x3_blocks):
        # get 2d slice (think of a plane in the x1 and x2 directions)
        slice_2d_df = get_tab_df_multicore(
                foldername, x1_blocks, x2_blocks, timestep=timestep, x3_block_offset=x3_index
                )

        if master_df is None:
            master_df = slice_2d_df
        else:
            master_df = master_df.append(slice_2d_df)

    return master_df

def compute_surface_dens(df):
    """
    Takes a dataframe containing the density at each point, returns
    a dataframe containing the surface density for each x
    """
    trim = df[["i-zone", "j-zone", "x1", "x2", "d", "dpar"]]
    #print(len(trim))
    imin = min(trim["i-zone"])
    imax = max(trim["i-zone"])

    dens = []
    for i in range(imin, imax+1):
        i_rows = trim.loc[(trim["i-zone"] == i)]
        row = [i, i_rows.iloc[0]["x1"]]

        dzarr = np.array([
            i_rows.iloc[j]["x2"] - i_rows.iloc[j-1]["x2"] for j in range(1,len(i_rows))
            ])
        par_surf_dens = sum(np.array(i_rows["dpar"])[1:] * dzarr)
        #print(surf_dens)
        row.append(par_surf_dens)

        #print(dzarr)
        #print(i_rows["d"])
        #print(np.array(i_rows["d"])[1:])
        gas_surf_dens = sum(np.array(i_rows["d"])[1:] * dzarr)
        row.append(gas_surf_dens)

        dens.append(row)

    densdf = pd.DataFrame(dens, columns=["i-zone", "x1", "dpar", "d"])

    return densdf

def integrate_df_3d(df, direction):
    """
    Integrate a 3d dataframe over an arbitrary direction
    honestly, this would probably be better if it turned the
    input df into a numpy cube, because i bet you could vectorize
    it and it would be faster, but whatever

    df : dataframe containing positions, dpar, etc.
    direction : one of "x1", "x2", "x3"
    """
    def direction_to_index(direction):
        return int(direction[1:]) - 1
    def direction_to_index_label(direction):
        return chr( direction_to_index(direction) + 105 )
    def get_direction_index_compl(dir_indx):
        return sorted([ (dir_index + n) % 3 for n in (1,2) ])
    def index_to_direction(dir_indx):
        return "x" + str(dir_indx + 1)

    dir_index = direction_to_index(direction)

    # directions orthogonal to integrating direction
    dir1_index, dir2_index = get_direction_index_compl(dir_index)
    dir1 = index_to_direction(dir1_index)
    dir2 = index_to_direction(dir2_index)
    dir1_index_label = direction_to_index_label(dir1) + "-zone"
    dir2_index_label = direction_to_index_label(dir2) + "-zone"

    dir1_min = min(df[dir1_index_label])
    dir1_max = max(df[dir1_index_label])
    dir2_min = min(df[dir2_index_label])
    dir2_max = max(df[dir2_index_label])

    new_rows = []
    for n in range(dir1_min, dir1_max+1):
        for m in range(dir2_min, dir2_max+1):
            locator = (df[dir1_index_label] == n) & (df[dir2_index_label] == m)
            nm_location_df = df.loc[locator]
            row = [n, m, nm_location_df.iloc[0][dir1], nm_location_df.iloc[0][dir2]]

            d_dir = np.abs(nm_location_df.iloc[1][direction] - nm_location_df.iloc[0][direction])
            compute_avgs = lambda arr : np.array([
                    0.5*(arr[l] + arr[l+1]) for l in range(0, len(arr)-1)
                    ])
            def compute_integral(col):
                vals = np.array(nm_location_df[col])
                avgs = compute_avgs(vals)
                return sum(avgs * d_dir)

            dpar_integral = compute_integral("dpar")
            d_integral = compute_integral("d")
            row.append(dpar_integral)
            row.append(d_integral)

            new_rows.append(row)

    columns = [dir1_index_label, dir2_index_label, dir1, dir2, "dpar", "d"]
    integrated_df = pd.DataFrame(new_rows, columns=columns)
    return integrated_df

def get_last_tab_time(folder):
    lsproc = subprocess.Popen(["ls", folder], stdout=subprocess.PIPE)
    grepproc = subprocess.Popen(["grep", "tab"], stdin=lsproc.stdout, stdout=subprocess.PIPE)
    tailproc = subprocess.Popen(["tail", "-1"], stdin=grepproc.stdout, stdout=subprocess.PIPE)

    stdout, stderr = tailproc.communicate()
    #print(stdout.decode("utf-8"))
    return int( stdout.decode("utf-8").split(".")[1] )

def load_surfdenscsv(filename):
    filevals = np.loadtxt(filename, delimiter=",")
    t_vals = filevals[:,0] # the first column
    surf_dens_arr = filevals[:,1:] # all but the first column
    return t_vals, surf_dens_arr

def get_df(filename):
    # i == x1 == r, j == x2 == z
    header_line = get_tab_header_line(filename)
    sep = "(?<!#)\s+"
    df = pd.read_csv(filename, sep=sep, header=header_line, engine="python")

    # column names in the file suck, so renaming them
    newcols = {}
    for col in list(df):
        newname = col[col.find("=")+1:]
        newcols[col] = newname
    df = df.rename(columns=newcols)
    df = df.apply(pd.to_numeric, errors="coerce")
    return df

def get_tab_header_line(filename):
    header_line = 0
    with open(filename, "r") as f:
        line = f.readline()
        while line != "":
            if line.find("[1]") != -1:
                break
            else:
                header_line += 1
                line = f.readline()
    return header_line

def notify(msg):
    notifystr = "notify-send -i /home/jeremy/Pictures/icons/python-icon-32.png Python '{0}'".format(msg)
    prc = subprocess.Popen(shlex.split(notifystr), stdout=subprocess.PIPE)

def Q(z, jmax):
    ans = 0
    for j in range(1, jmax+1):
        ans += (-1)**(j-1) * np.exp(-2 * z**2 * j**2)
    return 2*ans

def get_max_dist(cumdist):
    uniformline = np.array(range(0,len(cumdist)))/(len(cumdist)+0.0)
    return max(cumdist - uniformline)

def get_p(D, n):
    return Q(D*np.sqrt(n), 10000)


"""@namespace IMP.pmi.restraints.proteomics
Restraints for handling various kinds of proteomics data.
"""

from __future__ import print_function
import IMP
import IMP.core
import IMP.algebra
import IMP.atom
import IMP.container
import IMP.pmi
import IMP.pmi.tools
import IMP.pmi.output
import numpy
import math
import sys
import warnings

class ConnectivityRestraint(object):

    '''
    generate a connectivity restraint between domains
    setting up the restraint
    example:
    sel1 = IMP.atom.Selection(root_hier, molecule="Rpb3",
                              residue_indexes=range(1,100))
    sel2 = IMP.atom.Selection(root_hier, molecule="Rpb4",
                              residue_indexes=range(1,100))
    cr=restraints.ConnectivityRestraint((sel1, sel2), label='CR1')
    cr.add_to_model()
    Multistate support =No
    Resolution=Yes
    '''

    def __init__(self, domains, kappa=10.0, resolution=None, label="None"):
        self.weight = 1.0
        self.kappa = kappa
        self.label = label

        cr = IMP.atom.create_connectivity_restraint(
            domains, self.kappa, self.label)
        self.m = cr.get_model()
        self.rs = IMP.RestraintSet(self.m, label)
        self.rs.add_restraint(cr)

    def set_label(self, label):
        self.label = label
        self.rs.set_name(label)
        for r in self.rs.get_restraints():
            r.set_name(label)

    def add_to_model(self):
        IMP.pmi.tools.add_restraint_to_model(self.m, self.rs)

    def get_restraint(self):
        return self.rs

    def get_restraints(self):
        rlist = []
        for r in self.rs.get_restraints():
            rlist.append(IMP.core.PairRestraint.get_from(r))
        return rlist

    def set_weight(self, weight):
        self.weight = weight
        self.rs.set_weight(weight)

    def get_output(self):
        output = {}
        score = self.weight * self.rs.unprotected_evaluate(None)
        output["_TotalScore"] = str(score)
        output["ConnectivityRestraint_" + self.label] = str(score)
        return output

#

class CompositeRestraint(object):

    '''
    handleparticles a list of particles
    compositeparticles is a list of list of particles
    '''

    def __init__(self, handle_particles, composite_particles, cut_off=5.0,
                 lam=1.0, plateau=0.0, resolution=None, label="None"):
        # composite particles: all particles beside the handle
        self.label = label

        hs = IMP.pmi.tools.input_adaptor(handle_particles, resolution,
                                         flatten=True)
        self.handleparticles = [h.get_particle() for h in hs]
        self.m = self.handleparticles[0].get_model()
        self.rs = IMP.RestraintSet(self.m, 'cr')

        self.compositeparticles = []
        compositeparticle_list = []
        for cp in composite_particles:
            hs = IMP.pmi.tools.input_adaptor(cp, resolution, flatten=True)
            tmplist = [h.get_particle() for h in hs]
            compositeparticle_list.append(tmplist)
            self.compositeparticles += tmplist

        ln = IMP.pmi.CompositeRestraint(
            self.m, self.handleparticles, cut_off, lam, True, plateau)

        for ps in compositeparticle_list:
            # composite particles is a list of list of particles
            ln.add_composite_particle(ps)

        self.rs.add_restraint(ln)

    def set_label(self, label):
        self.label = label

    def get_handle_particles(self):
        return self.handleparticles

    def get_composite_particles(self):
        return self.compositeparticles

    def get_restraint(self):
        return self.rs

    def add_to_model(self):
        IMP.pmi.tools.add_restraint_to_model(self.m, self.rs)

    def get_output(self):
        output = {}
        score = self.rs.unprotected_evaluate(None)
        output["_TotalScore"] = str(score)
        output["CompositeRestraint_" + self.label] = str(score)
        return output

#
class AmbiguousCompositeRestraint(object):

    '''
    this restraint allows ambiguous cross-linking between multiple copies
    excluding between symmetric copies
    It allows name ambiguity
    '''

    def __init__(self, root_hier, restraints_file, cut_off=5.0, lam=1.0,
                 plateau=0.01, resolution=None, label="None"):
        self.weight = 1.0
        self.m = root_hier.get_model()
        self.rs = IMP.RestraintSet(self.m, 'data')
        self.label = "None"
        self.pairs = []

        self.outputlevel = "low"
        self.cut_off = cut_off
        self.lam = lam
        self.plateau = plateau

        fl = IMP.pmi.tools.open_file_or_inline_text(restraints_file)

        for line in fl:

            tokens = line.split()
            # skip character
            if (tokens[0] == "#"):
                continue
            r1 = int(tokens[2])
            c1 = tokens[0]
            r2 = int(tokens[3])
            c2 = tokens[1]

            ps1 = IMP.atom.Selection(root_hier, resolution=resolution,
                                     molecule=c1, residue_index=r1)
            ps1 = ps1.get_selected_particles()
            hrc1 = [p.get_name() for p in ps1]

            def nosym_subset(ps):
                return [p for p in ps if not IMP.pmi.Symmetric.get_is_setup(p)
                        or IMP.pmi.Symmetric(p).get_symmetric() == 0]
            ps1nosym = nosym_subset(ps1)
            hrc1nosym = [p.get_name() for p in ps1nosym]

            if len(ps1) == 0:
                warnings.warn(
                    "AmbiguousCompositeRestraint: residue %d of chain %s "
                    "is not there" % (r1, c1), IMP.pmi.StructureWarning)
                continue

            ps2 = IMP.atom.Selection(root_hier, resolution=resolution,
                                     molecule=c2, residue_index=r2)
            ps2 = ps2.get_selected_particles()
            hrc2 = [p.get_name() for p in ps2]
            ps2nosym = nosym_subset(ps2)
            hrc2nosym = [p.get_name() for p in ps2nosym]

            if len(ps2) == 0:
                warnings.warn(
                    "AmbiguousCompositeRestraint: residue %d of chain %s "
                    "is not there" % (r2, c2), IMP.pmi.StructureWarning)
                continue

            cr = IMP.pmi.CompositeRestraint(
                self.m, ps1nosym, self.cut_off, self.lam, True, self.plateau)
            cr.add_composite_particle(ps2)

            self.rs.add_restraint(cr)
            self.pairs.append(
                (ps1nosym,
                 hrc1nosym,
                 c1,
                 r1,
                 ps2,
                 hrc2,
                 c2,
                 r2,
                 cr))

            cr = IMP.pmi.CompositeRestraint(
                self.m, ps1, self.cut_off, self.lam, True, self.plateau)
            cr.add_composite_particle(ps2nosym)

            self.rs.add_restraint(cr)
            self.pairs.append(
                (ps1,
                 hrc1,
                 c1,
                 r1,
                 ps2nosym,
                 hrc2nosym,
                 c2,
                 r2,
                 cr))

    def plot_restraint(
        self,
        maxdist=100,
            npoints=100):

        p1 = IMP.Particle(self.m)
        p2 = IMP.Particle(self.m)
        d1 = IMP.core.XYZR.setup_particle(p1)
        d2 = IMP.core.XYZR.setup_particle(p2)
        cr = IMP.pmi.CompositeRestraint(
            self.m,
            [p1],
            self.cut_off,
            self.lam,
            True,
            self.plateau)
        cr.add_composite_particle([p2])
        dists = []
        scores = []
        for i in range(npoints):
            d2.set_coordinates(
                IMP.algebra.Vector3D(maxdist / npoints * float(i), 0, 0))
            dists.append(IMP.core.get_distance(d1, d2))
            scores.append(cr.unprotected_evaluate(None))
        IMP.pmi.output.plot_xy_data(dists, scores)

    def set_label(self, label):
        self.label = label
        self.rs.set_name(label)
        for r in self.rs.get_restraints():
            r.set_name(label)

    def add_to_model(self):
        IMP.pmi.tools.add_restraint_to_model(self.m, self.rs)

    def get_hierarchies(self):
        return self.prot

    def get_restraint_sets(self):
        return self.rs

    def get_restraint(self):
        return self.rs

    def set_output_level(self, level="low"):
        # this might be "low" or "high"
        self.outputlevel = level

    def set_weight(self, weight):
        self.weight = weight
        self.rs.set_weight(weight)

    def get_output(self):
        # content of the cross-link database pairs
        # self.pairs.append((p1,p2,dr,r1,c1,r2,c2))
        output = {}
        score = self.weight * self.rs.unprotected_evaluate(None)
        output["_TotalScore"] = str(score)
        output["AmbiguousCompositeRestraint_Score_" + self.label] = str(score)
        for n, p in enumerate(self.pairs):

            ps1 = p[0]
            hrc1 = p[1]
            c1 = p[2]
            r1 = p[3]
            ps2 = p[4]
            hrc2 = p[5]
            c2 = p[6]
            r2 = p[7]
            cr = p[8]
            for n1, p1 in enumerate(ps1):
                name1 = hrc1[n1]

                for n2, p2 in enumerate(ps2):
                    name2 = hrc2[n2]
                    d1 = IMP.core.XYZR(p1)
                    d2 = IMP.core.XYZR(p2)
                    label = str(r1) + ":" + name1 + "_" + str(r2) + ":" + name2
                    output["AmbiguousCompositeRestraint_Distance_" +
                           label] = str(IMP.core.get_distance(d1, d2))

            label = str(r1) + ":" + c1 + "_" + str(r2) + ":" + c2
            output["AmbiguousCompositeRestraint_Score_" +
                   label] = str(self.weight * cr.unprotected_evaluate(None))

        return output

#
class SimplifiedPEMAP(object):

    def __init__(self, root_hier, restraints_file, expdistance, strength,
                 resolution=None):
        self.m = root_hier.get_model()
        self.rs = IMP.RestraintSet(self.m, 'data')
        self.label = "None"
        self.pairs = []

        self.outputlevel = "low"
        self.expdistance = expdistance
        self.strength = strength

        fl = IMP.pmi.tools.open_file_or_inline_text(restraints_file)

        for line in fl:

            tokens = line.split()
            # skip character
            if (tokens[0] == "#"):
                continue
            r1 = int(tokens[2])
            c1 = tokens[0]
            r2 = int(tokens[3])
            c2 = tokens[1]
            pcc = float(tokens[4])

            ps1 = IMP.atom.Selection(root_hier, resolution=resolution,
                                     molecule=c1, residue_index=r1,
                                     copy_index=0)
            ps1 = ps1.get_selected_particles()
            if len(ps1) == 0:
                warnings.warn(
                    "SimplifiedPEMAP: residue %d of chain %s is not there "
                    "(w/ %d %s)" % (r1, c1, r2, c2), IMP.pmi.StructureWarning)
                continue
            if len(ps1) > 1:
                warnings.warn(
                    "SimplifiedPEMAP: residue %d of chain %s selected "
                    "multiple particles" % (r1, c1), IMP.pmi.StructureWarning)
                continue

            ps2 = IMP.atom.Selection(root_hier, resolution=resolution,
                                     molecule=c2, residue_index=r2,
                                     copy_index=0)
            ps2 = ps2.get_selected_particles()
            if len(ps2) == 0:
                warnings.warn(
                    "SimplifiedPEMAP: residue %d of chain %s is not there "
                    "(w/ %d %s)" % (r1, c1, r2, c2), IMP.pmi.StructureWarning)
                continue
            if len(ps2) > 1:
                warnings.warn(
                    "SimplifiedPEMAP: residue %d of chain %s selected "
                    "multiple particles" % (r2, c2), IMP.pmi.StructureWarning)
                continue

            p1 = ps1[0]
            p2 = ps2[0]

            # This is harmonic potential for the pE-MAP data
            upperdist = self.get_upper_bond(pcc)
            limit = 0.5 * self.strength * 15.0 ** 2 + 10.0
            hub = IMP.core.TruncatedHarmonicUpperBound(
                upperdist, self.strength, 15, limit)

            # This is harmonic for the X-link
            df = IMP.core.SphereDistancePairScore(hub)
            dr = IMP.core.PairRestraint(self.m, df, (p1, p2))
            self.rs.add_restraint(dr)
            self.pairs.append((p1, p2, dr, r1, c1, r2, c2))

            # Lower-bound restraint
            lowerdist = self.get_lower_bond(pcc)
            limit = 0.5 * self.strength * 15.0 ** 2 + 10.0
            hub2 = IMP.core.TruncatedHarmonicLowerBound(
                lowerdist, self.strength, 15, limit)

            # This is harmonic for the X-link
            df2 = IMP.core.SphereDistancePairScore(hub2)
            dr2 = IMP.core.PairRestraint(self.m, df2, (p1, p2))
            self.rs.add_restraint(dr2)
            self.pairs.append((p1, p2, dr2, r1, c1, r2, c2))

    def get_upper_bond(self, pearsoncc):
        # return (pearsoncc-1.)/-0.0075
        return (pearsoncc - .5) / (-0.005415)

    def get_lower_bond(self, pearsoncc):
        return (pearsoncc - 1.) / -0.0551

    def set_label(self, label):
        self.label = label

    def add_to_model(self):
        IMP.pmi.tools.add_restraint_to_model(self.m, self.rs)

    def get_hierarchies(self):
        return self.prot

    def get_restraint_sets(self):
        return self.rs

    def set_output_level(self, level="low"):
        # this might be "low" or "high"
        self.outputlevel = level

    def get_output(self):
        # content of the cross-link database pairs
        # self.pairs.append((p1,p2,dr,r1,c1,r2,c2))
        output = {}
        score = self.rs.unprotected_evaluate(None)
        output["_TotalScore"] = str(score)
        output["SimplifiedPEMAP_Score_" + self.label] = str(score)
        for i in range(len(self.pairs)):

            p0 = self.pairs[i][0]
            p1 = self.pairs[i][1]
            crosslinker = 'standard'
            ln = self.pairs[i][2]
            resid1 = self.pairs[i][3]
            chain1 = self.pairs[i][4]
            resid2 = self.pairs[i][5]
            chain2 = self.pairs[i][6]

            label = str(resid1) + ":" + chain1 + "_" + \
                str(resid2) + ":" + chain2
            output["SimplifiedPEMAP_Score_" + crosslinker + "_" +
                   label] = str(ln.unprotected_evaluate(None))

            d0 = IMP.core.XYZ(p0)
            d1 = IMP.core.XYZ(p1)
            output["SimplifiedPEMAP_Distance_" +
                   label] = str(IMP.core.get_distance(d0, d1))

        return output

class SetupConnectivityNetworkRestraint(object):

    '''
    generates and wraps a IMP.pmi.ConnectivityRestraint between domains
    example:
    cr=restraints.ConnectivityNetworkRestraint(
        simo,["CCC",(1,100,"TTT"),(100,150,"AAA")])
    cr.add_to_model()
    cr.set_label("CR1")

    Multistate support =No
    Selection type=selection tuple
    Resolution=Yes
    '''

    def __init__(self, objects, kappa=10.0, resolution=1.0, label="None"):

        self.weight = 1.0
        self.kappa = kappa
        self.label = label
        if self.label == "None":
            self.label = str(selection_tuples)  # noqa: F821

        hiers = []

        for obj in objects:
            hiers.append(IMP.pmi.tools.input_adaptor(
                obj, resolution, flatten=True))
            self.m = hiers[0][0].get_model()

        cr = ConnectivityNetworkRestraint(self.m)
        for hs in hiers:
            cr.add_particles([h.get_particle() for h in hs])
        self.rs = IMP.RestraintSet(self.m, label)
        self.rs.add_restraint(cr)

    def set_label(self, label):
        self.label = label
        self.rs.set_name(label)
        for r in self.rs.get_restraints():
            r.set_name(label)

    def add_to_model(self):
        IMP.pmi.tools.add_restraint_to_model(self.m, self.rs)

    def get_restraint(self):
        return self.rs

    def get_restraints(self):
        rlist = []
        for r in self.rs.get_restraints():
            rlist.append(IMP.core.PairRestraint.get_from(r))
        return rlist

    def set_weight(self, weight):
        self.weight = weight
        self.rs.set_weight(weight)

    def get_output(self):
        output = {}
        score = self.weight * self.rs.unprotected_evaluate(None)
        output["_TotalScore"] = str(score)
        output["ConnectivityNetworkRestraint_" + self.label] = str(score)
        return output

class ConnectivityNetworkRestraint(IMP.Restraint):

    '''
    a python restraint that computes the score for a composite of proteins
    Authors: G. Bouvier, R. Pellarin. Pasteur Institute.
    '''

    def __init__(self, m, slope=1.0, theta=0.0, plateau=0.0000000001,
                 linear_slope=0.015):
        '''
        input a list of particles, the slope and theta of the sigmoid potential
        theta is the cutoff distance for a protein-protein contact
        '''
        # Import networkx here so that we don't introduce it as a dependency
        # for *every* proteomics restraint, only this one
        import networkx
        self.networkx = networkx
        IMP.Restraint.__init__(self, m, "ConnectivityNetworkRestraint %1%")
        self.slope = slope
        self.theta = theta
        self.linear_slope = linear_slope
        self.plateau = plateau
        self.particles_blocks = []
        self.particle_list = []

    def get_number_of_particle_blocks(self):
        return len(self.particles_blocks)

    def get_number_of_particles_for_block(self, block_index):
        return len(self.particles_blocks[block_index])

    def add_particles(self, particles):
        self.particles_blocks.append(particles)
        self.particle_list += particles

    def get_full_graph(self):
        '''
        get the full graph of distances between every particle pair
        '''
        import scipy.spatial
        pdist_array = numpy.array(
            IMP.pmi.get_list_of_bipartite_minimum_sphere_distance(
                self.particles_blocks))
        pdist_mat = scipy.spatial.distance.squareform(pdist_array)
        pdist_mat[pdist_mat < 0] = 0
        graph = self.networkx.Graph(pdist_mat)
        return graph

    def get_minimum_spanning_tree(self):
        """
        return the minimum spanning tree
        """
        graph = self.get_full_graph()
        graph = self.networkx.minimum_spanning_tree(graph)
        return graph

    def sigmoid(self, x):
        '''
        a sigmoid function that scores the probability of a contact
        between two proteins
        '''
        # return 1 - (x)**self.slope/ float(((x)**self.slope +
        # self.theta**self.slope))
        argvalue = (x - self.theta) / self.slope
        return 1.0 - (1.0 - self.plateau) / (1.0 + math.exp(-argvalue))

    def unprotected_evaluate(self, da):
        graph = self.get_minimum_spanning_tree()
        score = 0.0
        for e in graph.edges():
            dist = graph.get_edge_data(*e)['weight']
            prob = self.sigmoid(dist)
            score += -numpy.log(prob)
            score += self.linear_slope * dist
        return score

    def do_get_inputs(self):
        return self.particle_list

class FuzzyBoolean(object):

    '''
    Fully Ambiguous Restraint that can be built using boolean logic
    R. Pellarin. Pasteur Institute.
    '''

    def __init__(self, p1, operator=None, p2=None):
        '''
        input a list of particles, the slope and theta of the sigmoid potential
        theta is the cutoff distance for a protein-protein contact
        '''
        if isinstance(p1, FuzzyBoolean) and isinstance(p2, FuzzyBoolean):
            self.operations = [p1, operator, p2]
            self.value = None
        else:
            self.operations = []
            self.value = p1

    def __or__(self, FuzzyBoolean2):
        return FuzzyBoolean(self, self.or_, FuzzyBoolean2)

    def __and__(self, FuzzyBoolean2):
        return FuzzyBoolean(self, self.and_, FuzzyBoolean2)

    def and_(self, a, b):
        return a * b

    def or_(self, a, b):
        return 1.0 - (1.0 - a) * (1.0 - b)

    def evaluate(self):

        if len(self.operations) == 0:
            return self.value
        FuzzyBoolean1, op, FuzzyBoolean2 = self.operations

        return op(FuzzyBoolean1.evaluate(), FuzzyBoolean2.evaluate())

class FuzzyRestraint(IMP.Restraint):

    '''
    Fully Ambiguous Restraint that can be built using boolean logic
    R. Pellarin. Pasteur Institute.
    '''
    plateau = 0.00000000000001
    theta = 5.0
    slope = 2.0
    innerslope = 0.01

    def __init__(self, m, p1, p2, operator=None):
        '''
        input a list of particles, the slope and theta of the sigmoid potential
        theta is the cutoff distance for a protein-protein contact
        '''
        IMP.Restraint.__init__(self, m, "FuzzyRestraint %1%")
        self.m = m
        self.min = sys.float_info.min
        if isinstance(p1, FuzzyRestraint) and isinstance(p2, FuzzyRestraint):
            self.operations = [p1, operator, p2]
            self.particle_pair = None
        elif isinstance(p1, FuzzyRestraint) and p2 is None:
            self.operations = [p1, operator, None]
            self.particle_pair = None
        else:
            self.operations = []
            self.particle_pair = (p1, p2)

    def __or__(self, FuzzyRestraint2):
        return FuzzyRestraint(self.m, self, FuzzyRestraint2, self.or_)

    def __and__(self, FuzzyRestraint2):
        return FuzzyRestraint(self.m, self, FuzzyRestraint2, self.and_)

    def __invert__(self):
        return FuzzyRestraint(self.m, self, None, self.invert_)

    def and_(self, a, b):
        c = a + b
        return c

    def or_(self, a, b):
        c = math.exp(-a) + math.exp(-b) - math.exp(-a - b)
        return -math.log(c)

    def invert_(self, a):
        c = 1.0 - math.exp(-a)
        return -math.log(c)

    def evaluate(self):
        if len(self.operations) == 0:
            return self.distance()
        FuzzyRestraint1, op, FuzzyRestraint2 = self.operations

        if FuzzyRestraint2 is not None:
            return op(FuzzyRestraint1.evaluate(), FuzzyRestraint2.evaluate())
        else:
            return op(FuzzyRestraint1.evaluate())

    def distance(self):
        d1 = IMP.core.XYZ(self.particle_pair[0])
        d2 = IMP.core.XYZ(self.particle_pair[1])
        d = IMP.core.get_distance(d1, d2)
        argvalue = (d-self.theta)/self.slope
        return (-math.log(1.0 - (1.0-self.plateau) / (1.0+math.exp(-argvalue)))
                + self.innerslope*d)

    def add_to_model(self):
        IMP.pmi.tools.add_restraint_to_model(self.m, self)

    def unprotected_evaluate(self, da):
        return self.evaluate()

    def __str__(self):
        if len(self.operations) == 0:
            return str(self.particle_pair)
        FuzzyRestraint1, op, FuzzyRestraint2 = self.operations
        if FuzzyRestraint2 is not None:
            return str(FuzzyRestraint1) + str(op) + str(FuzzyRestraint2)
        else:
            return str(FuzzyRestraint1) + str(op)

    def do_get_inputs(self):
        if len(self.operations) == 0:
            return list(self.particle_pair)
        FuzzyRestraint1, op, FuzzyRestraint2 = self.operations
        if FuzzyRestraint2 is not None:
            return list(set(FuzzyRestraint1.do_get_inputs()
                            + FuzzyRestraint2.do_get_inputs()))
        else:
            return list(set(FuzzyRestraint1.do_get_inputs()))


from google_auth_oauthlib.flow import InstalledAppFlow
from googleapiclient.discovery import build
from google.oauth2.credentials import Credentials
from google.auth.transport.requests import Request
import os

# All scopes together
ALL_SCOPES = [
    'https://www.googleapis.com/auth/contacts.readonly',
    'https://www.googleapis.com/auth/calendar.readonly',
    'https://www.googleapis.com/auth/gmail.readonly'
]

CREDENTIALS_FILE = 'env/token.json'

def get_credentials(scopes):
    creds = None

    # Load credentials from the file if it exists
    if os.path.exists(CREDENTIALS_FILE):
        creds = Credentials.from_authorized_user_file(CREDENTIALS_FILE, scopes)

    # Refresh or obtain new credentials
    if not creds or not creds.valid:
        if creds and creds.expired and creds.refresh_token:
            creds.refresh(Request())  # Use `Request()` as the required argument for `refresh()`
        else:
            flow = InstalledAppFlow.from_client_secrets_file('env/oauth2_credentials.json', scopes)
            creds = flow.run_local_server(port=0)

        # Save the credentials for the next run
        with open(CREDENTIALS_FILE, 'w') as token:
            token.write(creds.to_json())

    return creds

# Build the service using credentials
def get_authenticated_service(api, version, creds):
    return build(api, version, credentials=creds)

# Fetch contacts
def get_contacts(creds):
    service = get_authenticated_service('people', 'v1', creds)
    results = service.people().connections().list(
        resourceName='people/me',
        personFields='names,emailAddresses'
    ).execute()
    print("Data type of get_contacts results: ", type(results))
    # print("First 500 characters of results: ", json.dumps(results, indent=4)[:500])
    return results  # Add this line

# Fetch calendar events
def get_calendar_events(creds):
    service = get_authenticated_service('calendar', 'v3', creds)
    results = service.events().list(calendarId='primary').execute()
    print("Data type of get_calendar_events results: ", type(results))
    # print("First 500 characters of results: ", json.dumps(results, indent=4)[:500])
    return results  # Add this line

# Fetch emails
def get_emails(creds):
    service = get_authenticated_service('gmail', 'v1', creds)
    results = service.users().messages().list(userId='me', maxResults=10).execute()
    messages = results.get('messages', [])
    full_messages = []  # List to hold the full message details
    for message in messages:
        msg = service.users().messages().get(userId='me', id=message['id']).execute()
        full_messages.append(msg)
    return full_messages

if __name__ == "__main__":
    creds = get_credentials(ALL_SCOPES)
    get_contacts(creds)
    get_calendar_events(creds)
    get_emails(creds)


"""
This helps in finding the means and standards of the images to normalize before training.

  To run
python3 calculate_means_std.py -i path/to/image/folder/
"""

import argparse
import subprocess
import yaml
import os
import sys
sys.path.remove("/opt/ros/kinetic/lib/python2.7/dist-packages")
import cv2
import numpy as np

def is_image(filename):
  return any(filename.endswith(ext) for ext in ['.jpg', '.png'])

if __name__ == '__main__':
  parser = argparse.ArgumentParser()

  parser.add_argument(
      '--image', '-i',
      type=str,
      required=True,
      default=None,
      help='Directory to get the images from. If not passed, do from scratch!'
  )
  FLAGS, unparsed = parser.parse_known_args()

  # print summary of what we will do
  print("----------")
  print("INTERFACE:")
  #
  print("image dir", FLAGS.image)
  print("----------\n")

  print("----------\n")

  #

  # create list of images and examine their pixel values
  filenames = [os.path.join(dp, f) for dp, dn, fn in os.walk(
      os.path.expanduser(FLAGS.image)) for f in fn if is_image(f)]

  # examine individually pixel values
  counter = 0.0
  pix_val = np.zeros(3, dtype=np.float)
  for filename in filenames:
    # analize
    print("Accumulating mean", filename)

    # open as rgb
    cv_img = cv2.imread(filename, cv2.IMREAD_COLOR)
    cv_img = cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB)

    # normalize to 1
    cv_img = cv_img.astype(np.float) / 255.0

    # count pixels and add them to counter
    h, w, d = cv_img.shape
    counter += h * w

    # sum to moving pix value counter in each channel
    pix_val += np.sum(cv_img, (0, 1))

  # calculate means
  means = (pix_val / counter)

  # means
  print("means(rgb): ", means)

  # pass again and calculate variance
  pix_var = np.zeros(3, dtype=np.float)
  for filename in filenames:
    # analizel
    print("Accumulating variance", filename)

    # open as rgb
    cv_img = cv2.imread(filename, cv2.IMREAD_COLOR)
    cv_img = cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB)

    # normalize to 1
    cv_img = cv_img.astype(np.float) / 255.0

    # sum to moving pix value counter in each channel
    pix_var += np.sum(np.square(cv_img - means), (0, 1))

  # calculate the standard deviations
  stds = np.sqrt(pix_var / counter)
  print("stds(rgb): ", stds)

  # finalize by printing both
  print("*" * 80)
  print("means(rgb): ", means)
  print("stds(rgb): ", stds)
  print("*" * 80)


"""
words list가 링크만 있고, 글자 데이터가 없음. 글자 리스트를 받아오기 위한 파일
"""

from selenium import webdriver
from bs4 import BeautifulSoup
import pandas as pd

class Get_chndic_data:

    def __init__(self, link):
        self.link = link
        self.get_data = []

    def beautiful_soup(self, link):
        """
        Beautiful Soup Type의 객체를 return
        return값에 다음과 같은 메소드 가능
        data = soup.find('div', id='container').find('div', class_='section_hsk')
        :param link: https://zh.dict.naver.com/ 뒤에 들어갈 {letter_link}  주소
        :return: row 데이터
        """
        chrome_options = webdriver.ChromeOptions()
        chrome_options.add_argument('headless')
        chrome_options.add_argument('window-size=1920x1080')
        chrome_options.add_argument("disable-gpu")
        # https://beomi.github.io/2017/09/28/HowToMakeWebCrawler-Headless-Chrome/

        # driver = webdriver.Chrome("D:/dev/chromedriver.exe", chrome_options=chrome_options)  # 집에서 chromedriver 경로
        driver = webdriver.Chrome("C:/Users/user/Downloads/chromedriver.exe", chrome_options=chrome_options)
        # 학원에서 chromedriver 경로
        url = f'https://zh.dict.naver.com/{link}'
        driver.get(url)
        driver.minimize_window()
        content = driver.page_source.encode('utf-8').strip()
        soup = BeautifulSoup(content, "html.parser")
        driver.close()
        return soup

    def find_letter_inf(self):
        """
        글자 link를 이용해 글자와 글자 뜻, 병음, 단어일 경우 구성 글자와 구성 글자 링크를 받아옴
        :return:  위 자료로 구성된 list
        """
        # 병음 추가해야 함.
        soup = self.beautiful_soup(self.link)
        letter = soup.find('div', id='container').find('div', class_="section section_entry _section_entry") \
            .find('div', class_="entry_title _guide_lang").find('strong', class_='word').text
        return letter

temp_df = pd.read_csv('../csv/hsk_words_listed.csv', encoding='UTF-8')
hsk_words_link = temp_df.iloc[:, 1]
index = 500
get_data_list = []

########################################## 파일명 변환 ##################################
for i in range(3000, 3100):
    while True:
        try:
            get_data = Get_chndic_data(hsk_words_link[i])
            get_data_list.append(get_data.find_letter_inf())
            df = pd.DataFrame(get_data_list)
            print(df.tail())
            df.to_csv(f'../csv/letters_list{3000}.csv')
            break
        except AttributeError:
            print('try again')
            continue


import math
import random

# read data from file
def getData(fn):
    # get data
    file = open(fn, 'r')
    read = file.readlines()
    file.close

    states = int(read[0]) # number of states
    rewards = [] # reward data (also state data)
    moves = [] # move data

    for i in range(1, len(read)):
        row = read[i].split('\n')[0].split(' ')

        # append
        if i <= states:
            row[1] = int(row[1])
            rewards.append(row)
        else:
            row.append(0)
            moves.append(row)

    return (rewards, moves)

# lookup table -> return index
def lookup(table, coln, value):
    for i in range(len(table)):
        if table[i][coln] == value: return i
    return -1

# Q Learning
def QLearning_(rewards, moves, state, doRandom, times, lastModifiedTimes):
    # print
    symbol = rewards[state][0] # symbol of state
    print('')
    print('symbol of state: ' + symbol)
    for i in range(len(moves)): print(moves[i])

    # find all possible actions
    nextstate = []
    for i in range(len(moves)):
        if moves[i][0] == symbol: nextstate.append(moves[i])
    print('possible next moves: ' + str(nextstate))

    # decide next action
    nextActionIndex = -1 # index of next action
    # RANDOM
    if random.random() < doRandom:
        print('decide RANDOM')
        next_ = random.randint(0, len(nextstate)-1)
        state = lookup(rewards, 0, nextstate[next_][1]) # update state
        nextActionIndex = next_
    # BEST state
    else:
        print('decide BEST')
        bestVal = 0
        bestIndex = -1
        for i in range(len(nextstate)):
            if nextstate[i][2] > bestVal:
                bestVal = nextstate[i][2]
                bestIndex = i

        # best of next state (same value as best choice)
        # choose best next state randomly
        bestNextState = []
        for i in range(len(nextstate)):
            if nextstate[i][2] == bestVal:
                bestNextState.append(nextstate[i])
        selectBestNext = random.randint(0, len(bestNextState)-1)

        state = lookup(rewards, 0, bestNextState[selectBestNext][1]) # update state
        nextActionIndex = i

    # get immediate reward r(s, a)
    r = rewards[state][1]

    # find new state s' -> y*max(a')(Q(s', a'))
    symbol = rewards[state][0]
    maxReward = 0 # next action that has max reward
    for i in range(len(moves)):
        if moves[i][0] == symbol:
            if moves[i][2] > maxReward: maxReward = moves[i][2]

    # update state
    value = r + 0.5 * maxReward
    if nextstate[nextActionIndex][2] != value:
        lastModifiedTimes = times
        nextstate[nextActionIndex][2] = value
        print('update ' + str(nextstate[nextActionIndex][0]) + ' to '
            + str(nextstate[nextActionIndex][1]) + ': ' + str(value))

    print('symbol of state : ' + symbol)
    print('immediate reward: ' + str(r))
    print('max reward      : ' + str(maxReward))

    return (rewards, moves, state, times, lastModifiedTimes)

# do Q Learning
def QLearning(rewards, moves):
    # print
    print('Rewards:')
    for i in range(len(rewards)): print(rewards[i])

    times = 0
    lastModifiedTimes = 0 # times when action table was lastly changed
    while 1:
        # round info
        times += 1
        doRandom = 10/(times+10) # probability of searching randomly
        print('')
        print('----------------<<<< ROUND ' + str(times) + ' >>>>----------------')

        # do Q learning until the last state
        state = 0
        while state != len(rewards)-1:
            (rewards, moves, state, times, lastModifiedTimes) = QLearning_(rewards, moves, state, doRandom, times, lastModifiedTimes)

        # break when converged
        if lastModifiedTimes * 2 + 10 < times: break

if __name__ == '__main__':
    (rewards, moves) = getData('QLearning.txt')
    QLearning(rewards, moves)


import os
from google.cloud import storage

class GoogleStorageLoader():
    def __init__(self) -> None:
        """Start Google Cloud clint - could be used for uploading to storage
        """
        os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "../content/key-bucket.json"
        self.client = storage.Client()

    def upload_to_bucket(self, bucket_name, source_file, destination):
        """uploads file to the bucket

        Args:
            bucket_name (str): _description_
            source_file (str): _description_
            destination (str): _description_
        """
        bucket = self.client.bucket(bucket_name)
        blob = bucket.blob(destination)
        blob.upload_from_filename(source_file)


# coding: utf-8

# # Heat Diffusion in Soils
#
# This Jupyter Notebook gives an example how to implement a 1D heat diffusion model in Python.
#
# First we need to import the packages which we will be using:
#

# In[1]:

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

import CoupledHeatWaterFlowTHe as cfun
import MyTicToc as mt

sns.set()

## Main

# In[0:] Domain & Soil properties
nIN = 51
# soil profile until 15 meters depth
zIN = np.linspace(-2.0, 0, num=nIN).reshape(nIN, 1)
# nIN = np.shape(zIN)[0]
zN = np.zeros(nIN - 1).reshape(nIN - 1, 1)
zN[0, 0] = zIN[0, 0]
zN[1:nIN - 2, 0] = (zIN[1:nIN - 2, 0] + zIN[2:nIN - 1, 0]) / 2
zN[nIN - 2, 0] = zIN[nIN - 1]
nN = np.shape(zN)[0]

ii = np.arange(0, nN - 1)
dzN = (zN[ii + 1, 0] - zN[ii, 0]).reshape(nN - 1, 1)
dzIN = (zIN[1:, 0] - zIN[0:-1, 0]).reshape(nIN - 1, 1)

# collect model dimensions in a pandas series: mDim
mDim = {'zN' : zN,
        'zIN' : zIN,
        'dzN' : dzN,
        'dzIN' : dzIN,
        'nN' : nN,
        'nIN' : nIN
        }
mDim = pd.Series(mDim)

# ## Definition of material properties
# In this section of the code we define the material properties

# Soil Properties
# [J/(m3 K)] volumetric heat capacity of soil solids
zetaSol = 2.235e6
# [J/(m3 K)] volumetric heat capacity of water (Fredlund 2006)
zetaWat = 4.154e6

# rhoW = 1000  # [kg/m3] density of water
rhoS = 2650  # [kg/m3] density of solid phase
rhoB = 1700  # %[kg/m3] dry bulk density of soil
n = 1 - rhoB / rhoS  # [-] porosity of soil = saturated water content.
qCont = 0.75  # quartz content

# collect soil parameters in a pandas Series: sPar
sPar = {'vGA': np.ones(np.shape(zN)) * 1 / 2.0,  # alpha[1/m]
        'vGN': np.ones(np.shape(zN)) * 2.0,  # n[-]
        'vGM': np.ones(np.shape(zN)) * (1 - 1 / 2.0),  # m = 1-1/n[-]
        'thS': np.ones(np.shape(zN)) * 0.4,  # saturated water content
        'thR': np.ones(np.shape(zN)) * 0.03,  # residual water content
        'KSat': np.ones(np.shape(zN)) * 0.25,  # [m/day]
        'vGE': 0.5,  # power factor for Mualem-van Genuchten
        'Cv': 1.0e-8,  # compressibility of compact sand [1/Pa]
        'viscRef': cfun.ViscosityWaterT(283.15),
        'qCont': qCont, # quartz content
        }
sPar = pd.Series(sPar)

# In[1:] Definition of the Boundary Parameters

# Read meteodata
meteo_data  = pd.read_excel('WieringermeerData_Meteo.xlsx')
meteo_data['num_date'] = meteo_data['datetime'].astype(np.int64)/(1e9*3600*24)
meteo_data.set_index('datetime',inplace=True)

# set simulation time to numeric dates from boudary data...
t_range = meteo_data['num_date'][:-1]
taxis = meteo_data.index[:-1]

# collect boundary parameters in a named tuple boundpar...
def BndTTop(t, bPar):
    if np.size(t)==1:
        t = np.array([t])
    bndT = np.zeros(len(t))
    for ii in range(len(t)):
        xy, md_ind, t_ind = np.intersect1d(bPar.meteo_data['num_date'], np.ceil(t[ii]), return_indices=True)
        topT = bPar.meteo_data['temp'].iloc[md_ind].values

        bndT[ii] = 273.15 + topT

    return bndT

def BndqWatTop(t, bPar):
    if np.size(t)==1:
        t = np.array([t])
    qBnd = np.zeros(len(t))

    for ii in range(len(t)):
        xy, md_ind, t_ind = np.intersect1d(bPar.meteo_data['num_date'], np.ceil(t[ii]), return_indices=True)
        rf = bPar.meteo_data['rain_station'].iloc[md_ind].values

        qBnd[ii] = -rf
    return qBnd

bPar = {'topBndFuncHeat': BndTTop,
        'meteo_data': meteo_data,
        'topCond': 'Robin',
        'lambdaRobTop': 1e9,
        'lambdaRobBot': 0,
        'TBndBot': 273.15 + 10,
        'topBndFuncWat': BndqWatTop, #topBndFuncWat(t,bPar)
        'bottomTypeWat': 'Robin', # Robin condition or Gravity condition
        'kRobBotWat': 0.05,  # Robin resistance term for bottom
        'hwBotBnd': 1.0,  # pressure head at lower boundary
        }

bPar = pd.Series(bPar)

# In[3:] Define Initial Conditions
zRef = -1.0 # depth of water table
hwIni = zRef - zN

TIni = np.ones(np.shape(zN)) * (10.0 + 273.15)  # K

sVecIni = np.concatenate([hwIni, TIni], axis=0)

# Time Discretization
tOut = np.linspace(t_range[0],t_range[365],365*5)
#tplot = taxis[0:50]
nOut = np.shape(tOut)[0]

nOut = len(tOut)
# tOut = np.sort(np.hstack((tOut1, bTime)))  # time

# copy initial vector to hw0. Apply squeeze to compress it to one dimension

mt.tic()
int_result = cfun.IntegrateCHWF(tOut, sVecIni, sPar, mDim, bPar)
mt.toc()

hWSim = int_result.y[0:nN]
TSim = int_result.y[nN:2*nN]
thSim = cfun.thFun(hWSim,sPar)
qWSim = cfun.WatFlux(tOut,hWSim,TSim,sPar,mDim,bPar)
qHSim = cfun.HeatFlux(tOut, TSim, hWSim, sPar, mDim, bPar)

#mt.tic()
#TOutPic, hwOutPic = himp.HalfImplicitPicar(tOut2, hw0, T0, sPar, mDim, bPar, tPar)
#mt.toc()

sns.set()
plt.close('all')
fig1, ax1 = plt.subplots(figsize=(7, 4))
ii = np.arange(nN-1, 0, -10)
ax1.plot(tOut, TSim[ii,].T, '-')
ax1.set_title('Temperature (ODE)')
ax1.set_xlabel('time (days)')
ax1.set_ylabel('temperature [K]')
ax1.legend(zN[ii])

fig2, ax2 = plt.subplots(figsize=(7, 7))
jj = np.arange(0, nOut)
ax2.plot(TSim[:, jj], zN, '-')
ax2.set_title('Temperature vs. depth (ODE)')
ax2.set_ylabel('depth [m]')
ax2.set_xlabel('temperature [K]')

fig3, ax3 = plt.subplots(figsize=(7, 4))
# plot fluxes after 2nd output time (initial rate is extreme due to initial conditions)
ax3.plot(tOut, qHSim[ii,:].T, '-')
ax3.set_title('Heat Flux vs. depth (ODE)')
ax3.set_ylabel('depth [m]')
ax3.set_xlabel('temperature [J/m2]')
ax3.legend(zN[ii])

fig4, ax4 = plt.subplots(figsize=(7, 4))
# plot the pressure head for different depths as a function of time
# in this case we plot every 20th layer.
ax4.plot(tOut, hWSim[ii,:].T, '-')
ax4.set_ylabel('pressure head [m]')
ax4.set_xlabel('time [d]')

#plot pressure head as a function of depth. Here we plot every time step
fig5, ax5 = plt.subplots(figsize=(7, 7))
ax5.plot(hWSim, zN, '-')
ax5.grid(b=True)
ax5.set_xlabel('pressure head [m]')
ax5.set_ylabel('depth [m]')

# plt.savefig('myfig.png')

fig6, ax6 = plt.subplots(figsize=(7, 7))
ax6.plot(thSim, zN, '-')
ax6.grid(b=True)
ax6.set_xlabel('water content [-]')
ax6.set_ylabel('depth [m]')

fig7, ax7 = plt.subplots(figsize=(7, 4))
# plot the pressure head for different depths as a function of time
# in this case we plot every 20th layer.
ax7.plot(tOut, thSim[ii,:].T, '-')
ax7.set_ylabel('water content [-]')
ax7.set_xlabel('time [d]')
ax7.legend(zN[ii])

fig8, ax8 = plt.subplots(figsize=(7, 4))
# plot fluxes after 2nd output time (initial rate is extreme due to initial conditions)
ax8.plot(tOut, qWSim[ii,:].T, '-')
ax8.set_title('Water Flux ')
ax8.set_ylabel('depth [m]')
ax8.set_xlabel('water flow [m/d]')
ax8.legend(zN[ii])

fig1.savefig('./figures_scenarios/3_figure1.png')
fig2.savefig('./figures_scenarios/3_figure2.png')
fig3.savefig('./figures_scenarios/3_figure3.png')
fig4.savefig('./figures_scenarios/3_figure4.png')
fig5.savefig('./figures_scenarios/3_figure5.png')
fig6.savefig('./figures_scenarios/3_figure6.png')
fig7.savefig('./figures_scenarios/3_figure7.png')
fig8.savefig('./figures_scenarios/3_figure8.png')
# import shelve

# filename='/tmp/shelve.out'
# my_shelf = shelve.open(filename,'n') # 'n' for new

# for key in dir():
#     try:
#         my_shelf[key] = globals()[key]
#     except TypeError:
#         #
#         # __builtins__, my_shelf, and imported modules can not be shelved.
#         #
#         print('ERROR shelving: {0}'.format(key))
# my_shelf.close()