# ! /usr/bin/env python3

import random
import sys
from copy import deepcopy

class Card(object):
    def __init__(self, suit, number):
        if suit not in ("S", "C", "H", "D"):
            # S: spade, C: club, H: heart, D: Diamond
            raise ValueError("ERROR: suit of card is not correct: "
                             + str(suit))
        self.__suit = suit
        if number not in range(1, 14):
            raise ValueError("ERROR: number of card is not correct: "
                             + str(number))
        self.__number = number

    @property
    def card(self):
        return (self.__suit, self.__number)

    @property
    def suit(self):
        return self.__suit

    @property
    def number(self):
        return self.__number

class Status(object):
    def __init__(self, money):
        self.__money = money
        self.__bet_money = 0
        if self.__money > 0:
            self.in_game = True
            self.alive = True
        else:
            self.in_game = False
            self.alive = False

    def bet(self, money):
        self.__money -= money
        self.__bet_money += money
        return money

class Dealer(object):
    def __init__(self, game_inst, players_input):
        # FIXED PARAMETERS
        self.__MIN_NUMBER_CARDS = 1  # smallest number of playing cards
        self.__MAX_NUMBER_CARDS = 13  # largest number of playing cards
        self.__SUITE = ['S', 'C', 'H', 'D']  # suit of playing cards
        self.__NUM_HAND = 2  # number of hands
        self.__NUM_MAX_FIELD = 5  # maximum number of field

        self.__game_inst = game_inst
        self.__players = deepcopy(players_input)  # instance of players
        self.__num_players = len(self.__players)  # number of players
        self.__num_handling_cards \
            = self.__NUM_HAND * self.__num_players + self.__NUM_MAX_FIELD
        # player's hand money list
        self.__money_each_player = self.__game_inst.accounts
        # create list of players' status
        self.__list_status = [Status(_money) for _money in self.__game_inst.accounts]
        print(self.__money_each_player)
        # cards list on the field
        self.__field = []
        # get players know dealer's instance
        for player in self.__players:
            player.get_know_dealer(self)
        # DB, SB, BB positioning
        self.__DB = self.__game_inst.DB
        self.__SB = self.__next_alive_player(self.__DB)
        self.__BB = self.__next_alive_player(self.__SB)
        self.__betting_cost = game_inst.minimum_bet  # betting money at first
        # player can raise betting money the multiple of game_inst.minimum_bet
        self.minimum_bet = game_inst.minimum_bet
        self.playercheck = [True]*self.__num_players  # 返答を毎度更新し、降りた時にFalseにする
        self.active_players_list = self.__players  # the list of actionable players
        self.bettingrate = [0]*self.__num_players  # 各々が賭けたお金を記録するリスト
        self.bettingrate[self.__SB] = 1  # small-blined bet 1$ at first
        self.bettingrate[self.__BB] = 2  # big-blined bet 2$ at first
        self.__pot = self.__list_status[self.__SB].bet(self.minimum_bet/2)
        self.__pot = self.__list_status[self.__BB].bet(self.minimum_bet)
        # this flag is used to compair to big-blined position
        self.flag_atfirst = 0
        # this flag is used to check the count of "call" and "fold"
        self.flag = 1
        # list of players respond("call"/"fold"/number of raise)
        self.resplist = []
        for i in range(self.__num_players):
            if self.__money_each_player[i] <= 0:
                self.playercheck[i] = False
                self.__list_status[i].in_game = False
                # player who have no money can't play new game
        # print the player of small-blined position
        print("SB Player", self.__SB)
        # print the player of big-blined position
        print("BB Player", self.__BB)

    # give ith as int, return next player's index who is in the game
    def __next_alive_player(self, ith):
        _i = (ith + 1) % self.__num_players
        while self.__list_status[_i].alive is False:
            _i = (_i + 1) % self.__num_players
        return _i

    # create a deck
    def __create_all_cards_stack(self):  # create list of [S1, S2, ..., D13]
        _cards = []
        for inumber in range(self.__MIN_NUMBER_CARDS, self.__MAX_NUMBER_CARDS+1):
            for suit in self.__SUITE:
                _cards.append(Card(suit, inumber))
        return _cards

    # handout cards to each player
    def handout_cards(self):
        self.__field = []
        self.__all_cards = self.__create_all_cards_stack()
        self.__handling_cards = random.sample(self.__all_cards, self.__num_handling_cards)
        self.__players_cards = []  # each player's hand
        for player in self.__players:
            self.__players_cards.append(
                    [self.__handling_cards.pop(i) for i in range(self.__NUM_HAND)]
                    )
            player.get_hand(self.__players_cards[-1])

    # open one card to a table
    def put_field(self):
        self.__field.append(self.__handling_cards.pop(0))

    # //////////////////////////////////////////////////////////////////////////
    # /////////////////ここから先get_responses関連///////////////////////////////
    # //////////////////////////////////////////////////////////////////////////
    # ask players what they want to do "fold, call, raise"

    # number of players who play new game
    def __num_possible_players(self):
        return len([i for i in self.__list_status if i.alive])

#    def SB_kosin(self):  # number of small-blined(0~3)
    def DB_update(self):  # number of small-blined(0~3)
        return self.__next_alive_player(self.__DB)

    # ＜出てくるリストや変数＞
    # self.resplistはplayerの返答をリストにしたもの
    # "call"/"fold"/int() 参加資格が無いあるいは訊き始める条件にはない場合はskipし"----"を格納する
    # self.player_numberはリストの何番目のplayerなのかを表したもの　追加した返答をルールに従うように補正する際関数に渡す
    # self.flag self.flag_atfirstはそれぞれcallかfoldが続いたカウント、1ターン目のBB2ターン目以降のSBまでskipするためのカウンター
    #
    # <関数の構成>
    # playerの番号を得る
    # 必要があればskip("----"格納)してそれ以外はplayerに返答を訊きリストに格納
    # 格納した返答がルールに従ったものになるように修正する
    # 最後にフラグと生きているplayerのリストを更新する
    def get_response_from_one_person(self, player):
        self.player_number = len(self.resplist)  # know the number of player 0~4
        if self.flag >= self.__num_players or len(self.active_players_list) == 1:
            self.resplist.append("----")  # skip players after fill the conditions to move nexat turn
        elif self.flag_atfirst <= self.__BB and self.__BB != self.player_number - 1:
            self.resplist.append("----")  # skip untill BB at first turn
        elif self.playercheck[self.player_number] is False:
            self.resplist.append("----")  # skip the player
        else:
            self.resplist.append(player.respond())  # get response from player
        self.hentounohosei(self.player_number)  # correct response to follow the game rules
        self.flagnokosin(self.player_number)  # move flags
        self.active_players()  # renew active_plyers_list

    def hentounohosei(self, i):
        # while文でflagがプレイヤー数になるという次の工程に移行する条件を定義
        # flagでレイズから次にレイズがあるまでカウントししている
        if self.flag_atfirst <= self.__BB and self.__BB != self.player_number - 1:
            self.flag = self.flag-1
        elif self.resplist[i] == "----":
            pass
        elif self.resplist[i] == "fold":
            self.playercheck[i] = False  # 降りる
            self.__list_status[i].in_game = False
        elif self.__money_each_player[i] <= self.__betting_cost:
            self.resplist[i] = "call"  # 掛け金に満たない場合で降りてないなら必然的にcall
            self.bettingrate[i] = self.__money_each_player[i]
        elif self.resplist[i] == "call" or 0:  # お金あるときのcall
            self.bettingrate[i] = self.__betting_cost
        elif type(self.resplist[i]) is str:
            self.resplist[i] = "call"
            self.bettingrate[i] = self.__betting_cost
        elif self.__betting_cost+self.minimum_bet >= self.__money_each_player[i]:
            self.bettingrate[i] = self.__money_each_player[i]
            self.flag = 0
            self.__betting_cost = self.__betting_cost+self.minimum_bet
        else:
            if self.minimum_bet > self.resplist[i]:
                # minimum_betより小さい金額ならminimum_betに修正
                self.resplist[i] = self.minimum_bet
                self.__betting_cost = self.__betting_cost+self.resplist[i]
                # call金額の更新
                self.bettingrate[i] = self.__betting_cost
            else:
                # minimum_betの整数倍をレイズするように返値を修正
                j = int(self.resplist[i]/self.minimum_bet)
                if self.__money_each_player[i] <= self.__betting_cost+self.minimum_bet*j:
                    j = int((self.__money_each_player[i]-self.__betting_cost)/self.minimum_bet)+1
                self.minimum_bet = self.minimum_bet*j
                self.resplist[i] = self.minimum_bet
                # minimum_betの更新
                self.__betting_cost = self.__betting_cost+self.resplist[i]
                # call金額の更新
                self.bettingrate[i] = self.__betting_cost
            self.flag = 0

    def kakekinhosei(self):
        for i in range(0, self.__num_players):  # 最終的な掛け金の補正
            if self.__money_each_player[i] <= self.bettingrate[i]:
                self.bettingrate[i] = self.__money_each_player[i]

    def flagnokosin(self, i):
        self.flag = self.flag+1
        self.flag_atfirst = self.flag_atfirst+1

    def active_players(self):
        self.active_players_list = []
        for j in range(self.__num_players):  # 降りなかった人をリストで返す
            if self.playercheck[j] is True:
                self.active_players_list.append(self.__players[j])

    def printingdate(self):
        print("next_turn_players_list", [i.__class__.__name__ for i in self.active_players_list])
        # 次のターン参加する人のリスト
        print("betting_rate", self.__betting_cost)  # レイズを繰り返した最終的にcallがそろった時の金額
        print("personal_betting_money", self.bettingrate)
        # 降りた人も含めてこの時点でいくら賭けたかのリスト
        print()
        print()
        if len(self.field) == 5:
            print("--------------------------------------------")
            for i in range(0, self.__num_players):
                self.__money_each_player[i] = self.__money_each_player[i]-self.bettingrate[i]
            print("hanteimae-no-syozikin = ", self.__money_each_player)
            print("syousya-hantei-taisyousya = ", [i.__class__.__name__ for i in self.active_players_list])
            self.pot = sum(self.bettingrate)
            print("pot = ", self.pot)

    def get_responses(self):  # playersから返事を次のターンに進められるまで聞き続ける
        if len(self.field) != 0:
            self.flag = 0
            self.__BB = (self.__SB-1) % self.__num_players
        self.flag_atfirst = 0
        while self.flag < self.__num_players and len(self.active_players_list) != 1:
            # while文でflagがプレイヤー数になるという次の工程に移行する条件を定義
            if len(self.resplist) == self.__num_players:
                self.resplist = []
            # 1人ずつ聞いて補正して反映させる
            self.resp = [
                self.get_response_from_one_person(player) for player in self.__players
                        ]
            self.kakekinhosei()  # 持ち金を超えた掛け金の補正
            print(self.resplist)
        self.printingdate()  # 必要なデータをprint
        # ///////////////////////////////////////////////////////////////////////
        # /////////////////ここまでget_responses()関連////////////////////////////
        # ///////////////////////////////////////////////////////////////////////

    def calc(self):
        self.active_players()
        winner = []
        winner_num = []
        i = 0
        j = 0
        winners_cards_list = []
        winner_score = 0
        roll = []
        for player in self.__players:
            if self.playercheck[i] is True:
                seven_cards = self.__players_cards[self.__players.index(player)] + self.field
                roll.append(self.calc_hand_score(seven_cards)[0])
                if winner_score < roll[j]:
                    winner = [self.active_players_list[j]]
                    winner_num = []
                    winner_num.append(i)
                    winner_score = roll[j]
                    winners_cards_list = [[i, seven_cards]]
                elif winner_score == roll[j]:
                    winner_num.append(i)
                    winner.append(self.active_players_list[j])
                    winners_cards_list.append([i, seven_cards])
                print()
                j = j+1
            i = i+1
        print("--------------------------------------------")
        print([i.__class__.__name__ for i in self.active_players_list], " = ", roll)  # print finalist and their score
        '''
        if  len(winner) != 0:
            winners_cards_list = self.deside_winner_from_highcard(winners_cards_list, winner_score)
            # 同じ役の人たちを比較して新しいwinners_card_kistを返す関数を作成する 引数は[[player番号, カードリスト], [player番号, カードリスト]...], 役のスコア]
        '''
        print("///////////////////////////////////////////////")
        print("/////////////winner", [i.__class__.__name__ for i in winner], "////////////////")  # print winner
        print("///////////////////////////////////////////////")
        winning_money = int(self.pot/len(winner))
        for i in range(len(winner)):
            self.__money_each_player[winners_cards_list[i][0]] = self.__money_each_player[winners_cards_list[i][0]] + winning_money

    # calculate best score from given set of cards
    # 担当：白井．7枚のカードリストを受け取り，役とベストカードを返します．
    def calc_hand_score(self, cards):  # 7カードリストクラスをもらう
        SS = ['S', 'C', 'H', 'D']
        suit_list = [0, 0, 0, 0]
        rtCrads = []

        (num, suit, card_list) = self.choice(cards)  # クラスからnum, suit, cardを抜き出す
        # card_list = sorted(card_list, key = lambda x: x[1]) # # <<DEBUG MODE>>
        # print("card:", card_list)
        pp = self.checkpair(cards)  # Kawadaさんの4cardsとか抜き出してリストにするやつ
        #  # REPLACE 1-->14
        num = self.rpc1(num)
        num.sort()
        nc = self.rpcards1(card_list)
        card_list = nc
        card_list = sorted(card_list, key=lambda x: x[1])  # 2ndでsort

        flash = 0
        straight = 0
        straight_flash = 0
        # for flash:make flash_list
        for SUIT in SS:
            if suit.count(SUIT) >= 5:  # flash
                flash = 1
                flash_list = []
                for i in range(len(card_list)):  # flashの数字だけ取り出す
                    if card_list[6-i][0] == SUIT:
                        flash_list.append(card_list[6-i])
        (straight, straight_list) = self.stlist(card_list)
        if straight == 1 and flash == 1:
            (st, st_list) = self.stlist(flash_list)
            if st == 1:
                score = 8
                straight_flash = 1

        # == JUDGE BELOW ==
        rtCards = []
        # Straight-Flash
        if straight_flash == 1:
            rtCards = st_list
        # 4cards
        elif pp[0] >= 1:
            score = 7
            for i in range(self.__MAX_NUMBER_CARDS):
                if num.count(14-i) == 4:
                    for n in range(len(card_list)):
                        if card_list[6-n][1] == 14-i:
                            rtCards.append(card_list[6-n])
                            card_list.remove(card_list[6-n])
            rtCards.append(card_list[2])
        # Fullhouse
        elif pp[1] == 2:  # 3c *2
            score = 6
            c = 0
            for i in range(self.__MAX_NUMBER_CARDS):
                if num.count(14-i) == 3:
                    num.remove(14-i)
                    num.remove(14-i)
                    num.remove(14-i)
                    for n in range(len(card_list)):
                        if card_list[6-n][1] == (14-i):
                            rtCards.append(card_list[6-n])
                            card_list.remove(card_list[6-n])
                    break
            for i in range(self.__MAX_NUMBER_CARDS):
                if num.count(14-i) == 3:
                    for n in range(len(card_list)):
                        if card_list[n][1] == (14-i):
                            rtCards.append(card_list[3-n])
                            card_list.remove(card_list[3-n])
                            c += 1
                            if c == 2:  # 小さい方の3cは2個だけ取る
                                break
        elif pp[1] == 1 and pp[2] >= 1:  # 3c+pair
            score = 6
            for i in range(self.__MAX_NUMBER_CARDS):
                if num.count(14-i) == 3:
                    num.remove(14-i)
                    num.remove(14-i)
                    num.remove(14-i)
                    for n in range(len(card_list)):
                        if card_list[6-n][1] == (14-i):
                            rtCards.append(card_list[6-n])
                            card_list.remove(card_list[6-n])
                    break
            for i in range(self.__MAX_NUMBER_CARDS):
                if num.count(14-i) == 2:
                    for n in range(len(card_list)):
                        if card_list[3-n][1] == (14-i):
                            rtCards.append(card_list[3-n])
                            card_list.remove(card_list[3-n])
                    break
        # Flash
        elif flash == 1:
            score = 5
            rtCards = flash_list[len(flash_list)-5:len(flash_list)]
        # Straight
        elif straight == 1:
            score = 4
            rtCards = straight_list
        # 3cards
        elif pp[1] == 1:
            score = 3
            for i in range(self.__MAX_NUMBER_CARDS):
                if num.count(14-i) == 3:
                    for n in range(len(card_list)):
                        if card_list[6-n][1] == (14-i):
                            rtCards.append(card_list[6-n])
                            card_list.remove(card_list[6-n])
                    break
            for i in range(2):
                rtCards.append(card_list[3-i])
                card_list.remove(card_list[3-i])
        # 2pairs
        elif pp[2] >= 2:
            score = 2
            c = 0
            for i in range(self.__MAX_NUMBER_CARDS):
                if num.count(14-i) == 2 and c != 2:
                    c += 1
                    num.remove(14-i)
                    num.remove(14-i)
                    for n in range(len(card_list)):
                        if card_list[len(card_list)-1-n][1] == (14-i):
                            rtCards.append(card_list[len(card_list)-1-n])
                            rtCards.append(card_list[len(card_list)-2-n])
                            card_list.remove(card_list[len(card_list)-1-n])  # 可変だからこうなる
                            card_list.remove(card_list[len(card_list)-1-n])
                            break
            rtCards.append(card_list[2])
        # 1pair
        elif pp[2] == 1:
            score = 1
            for i in range(self.__MAX_NUMBER_CARDS):
                if num.count(14-i) == 2:
                    for n in range(len(card_list)):
                        if card_list[6-n][1] == 14-i:
                            rtCards.append(card_list[6-n])
                            card_list.remove(card_list[6-n])
                    break
            for i in range(3):
                rtCards.append(card_list[4-i])
                card_list.remove(card_list[4-i])
        # no pair
        else:
            score = 0
            rtCards = card_list[2:7]

        # RETURN!!
        nc = self.rpcards2(rtCards)
        rtCards = nc
        return (score, rtCards)

    # for: calc_hand_score
    def choice(self, card_list):  # suit, num, cardのみを取り出してリスト化
        SS = ['S', 'C', 'H', 'D']
        suit = [0]*len(card_list)
        num = [0]*len(card_list)
        card = [0]*len(card_list)
        for i in range(len(card_list)):
            num[i] = card_list[i].number
            suit[i] = card_list[i].suit
            card[i] = card_list[i].card
        return (num, suit, card)

    # for straight:make straight_list
    def stlist(self, card_list):
        card_list = sorted(card_list, key=lambda x: x[1])
        num = [0]*len(card_list)
        for i in range(len(card_list)):
            num[i] = card_list[i][1]
        straight_list = []
        straight = 0
        num_list = [0]*15
        for card in num:  # 数字の個数カウント
            num_list[card] += 1
        for i in range(10):
            prod = num_list[14-i]*num_list[self.__MAX_NUMBER_CARDS-i]*num_list[12-i]*num_list[11-i]*num_list[10-i]
            if prod >= 1:
                straight = 1  # st宣言
                k = 0
                for t in range(len(card_list)):
                    if card_list[len(card_list)-1-t][1] == 14-i-k and k < 5:
                        straight_list.append(card_list[len(card_list)-1-t])
                        k += 1
                break
        return (straight, straight_list)

    # Kawadaさんのもの
    def checkpair(self, any_cards):  # ペアの評価方法
        pair = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]  # A~Kまでの13個のリスト要素を用意
        for i in range(0, len(any_cards)):  # カードの枚数ぶんだけ試行
            pair[any_cards[i].number-1] = pair[any_cards[i].number-1]+1  # カードのnumber要素を参照し先ほどのリストpairの対応要素のカウントを1つ増やす
        pairs = [0, 0, 0]  # pairsは[4カード有無, 3カードの有無, ペアの数]のリスト
        for i in range(0, self.__MAX_NUMBER_CARDS):  # pairの要素A~13すべて順に参照
            if pair[i] == 4:  # lその要素が４枚あるときpairs[0]のカウントを増やす
                pairs[0] = pairs[0]+1
            elif pair[i] == 3:  # 同様に3枚
                pairs[1] = pairs[1]+1
            elif pair[i] == 2:  # 同様に2枚
                pairs[2] = pairs[2]+1
        return pairs  # pairsは[4カード有無, 3カードの有無, ペアの数]のリスト

    def rpc1(self, cards):  # 最初に1-->14にする方 引数はリスト
        rp = []
        for card in cards:
            card = (card+11) % 13 + 2
            rp.append(card)
        return rp

    def rpc2(self, cards):  # 最後に14-->1に戻す方 引数はリスト
        rp = []
        for card in cards:
            card = (card-1) % 13 + 1
            rp.append(card)
        return rp

    def rpcards1(self, cards):  # 最初に1-->14にする方 引数はカードタプルリスト
        nc = []
        for i in range(len(cards)):
            ss = cards[i][0]
            nn = (cards[i][1]+11) % 13 + 2
            nc.append((ss, nn))
        return nc

    def rpcards2(self, cards):  # 最後に14-->1に戻す方 引数はカードタプルリスト
        nc = []
        for i in range(len(cards)):
            ss = cards[i][0]
            nn = (cards[i][1]-1) % 13 + 1
            nc.append((ss, nn))
        return nc

    # == for: calc_hand_score
    def judge_flash(self, cl1, cl2):  # == FLASH判定 shirai
        num1 = []
        num2 = []
        for k in range(5):
            num1.append(cl1[k][1])
            num2.append(cl2[k][1])
        self.rpc1(num1)
        self.rpc1(num2)
        for i in range(5):
            if max(num1) > max(num2):
                sc = 0
                break
            elif max(num1) < max(num2):
                sc = 1
                break
            elif max(num1) == max(num2):
                sc = 2
                num1.remove(max(num1))
                num2.remove(max(num2))
        return sc

    def judge_straight(self, cl1, cl2):  # == STRAIGHT判定 shirai
        num1 = []
        num2 = []
        for i in range(5):
            num1.append(cl1[i][1])
            num2.append(cl2[i][1])
        num1 = self.rpc1(num1)
        num2 = self.rpc1(num2)
        if max(num1) > max(num2):
            sc = 0
        elif max(num1) < max(num2):
            sc = 1
        elif max(num1) == max(num2):
            sc = 2
        return sc

    @property
    def field(self):
        return self.__field

    @property
    def list_of_players(self):
        return [i.__class__.__name__ for i in self.__players]

    @property
    def list_of_money(self):
        return self.__money_each_player

    def get_position(self,_your_inst):
        return self.__players.index(_your_inst)

    @property
    def DB(self):
        return self.__DB

    @property
    def betting_cost(self):
        return self.__betting_cost

    # obsolete #
    @property
    def money(self):
        return self.__betting_cost


import covid
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import datetime as dt
from matplotlib.font_manager import FontProperties
import pytz  # Add this import to handle time zones

goth = FontProperties(fname="C:/Windows/Fonts/msgothic.ttc")

optitle = {
    "xycoords": "data",
    "textcoords": "offset points",
    "fontproperties": goth,
    "size": 12
}

opplot = {
    "linewidth": 1.2,
    "alpha": 0.7,
    "color": "blue"
}

fpmi = {
    "Week": [dt.date(2023,5, 8), dt.date(2023, 5, 15), dt.date(2023, 5, 22), dt.date(2023, 5, 29), dt.date(2023, 6, 5), dt.date(2023, 6, 12), dt.date(2023, 6, 19), dt.date(2023, 6, 26)],
    "東京都": [2.40, 3.53, 3.96, 5.29],
    "滋賀県": [1.82, 2.07, 1.77, 2.47],
    "鳥取県": [2.69, 3.24, 2.86, 4.24],
    "全国": [2.63, 3.55, 3.63, 4.55],
    "大阪府":[1.79, 2.37, 2.75, 3.33, 4.33, 4.55, 5.16, 5.93]
}

def ex06(pref, odays, fac, epoch0, ndays, xlim, xticks, ylim, yticks, pngfile, log=False, rotation=28, pad=-2):
    filename = "newly_confirmed_cases_daily.csv"
    daily = covid.readCSV(filename)
    weekly1, weekly2 = covid.readXLSXfp()
    fig = plt.figure(figsize=(6, 5))
    ax = plt.axes([0.09, 0.11, 0.89, 0.83])
    xlim, xticks = [epoch0, epoch0 + dt.timedelta(days=ndays)], [epoch0 + dt.timedelta(days=i) for i in range(7, ndays, 4 * 7)]
    if log:
        ax.set_yscale("log")
    p = covid.prefecture[pref]
    covid.drawFrame("", optitle, ax, xlim, xticks, ylim, yticks, rotation=rotation, pad=pad)
    ax.set_title("大阪府 定点あたりの報告数の日別陽性者数への変換", fontproperties=goth, size=12, pad=4)
    ax.text(0, 1.05, "Y210231 北村拓也", transform=ax.transAxes, fontproperties=goth, size=9)

    weekly_pref = weekly2[pref]  # 特定の都道府県のデータを抽出

    # Convert dates to pandas Timestamp objects and specify time zone (Japan Standard Time)
    dates = [pd.Timestamp(date, tz=pytz.timezone("Asia/Tokyo")) for date in fpmi["Week"]]

    ax.plot(dates, fac * np.array(fpmi[pref]), label=f"{pref}FPMI × {fac}", **{**opplot, "color": "orange"})
    ax.legend(prop=goth, loc="upper right", bbox_to_anchor=(1.0, 1.0))
    plt.show()
    fig.savefig(pngfile, dpi=600)

if __name__ == "__main__":
    epoch0, ndays = dt.datetime(2022, 1, 10), 82 * 7
    xlim, xticks = [epoch0, epoch0 + dt.timedelta(days=ndays)], [epoch0 + dt.timedelta(days=i) for i in range(7, ndays, 4 * 7)]
    pref, odays, fac, ylim, yticks = "大阪府", 3, 350, [0.2, 200], [1, 10, 100, 1000, 10000, 100000]
    ex06(pref, odays, fac, epoch0, ndays, xlim, xticks, ylim, yticks, "ex05.png", log=True)


import bz2
import json
import spacy
from profanity_filter import ProfanityFilter
import numpy as np

nlp = spacy.load('en')
profanity_filter = ProfanityFilter(nlps={'en': nlp})
nlp.add_pipe(profanity_filter.spacy_component, last=True)

def construct_dict(files):
    lines = []
    search_dict = {}
    for filename in files:
        with bz2.open(filename, "rt") as bzinput:
            for i, line in enumerate(bzinput):
                # if i == 100: break
                tweets = json.loads(line)
                if "text" not in tweets:
                    continue
                if tweets["lang"] != "en":
                    continue
                if tweets["retweeted"]:
                    continue
                lines.append(tweets)

                words = line["text"].split()
                for word in words:
                    w = word.lower()
                    if w in search_dict:
                        search_dict[w].append(line["id_str"])
                    else:
                        search_dict[w] = [line["id_str"]]
    return lines, search_dict

def get_tweets(tgt, lines):
  for line in lines:
    if "text" not in line:
      continue
    if line["lang"] != "en":
      continue
    else:
      if line["id_str"] in tgt:
        yield line["text"]

def get_samples(new_term, dict, lines):

    if new_term in dict:
        ids = dict[new_term]
    else:
        return None

    tweets = get_tweets(ids, lines)

    filtered = []
    for tweet in tweets:
        doc = nlp(tweet)
        if not doc._.is_profane:
            filtered.append(tweets)

    return filtered

def sample_defs(def_list, new_term, num_samples= 4):

    samples = []

    while len(samples) <= num_samples:
        sample = np.random.sample(def_list)
        if sample["word"] == new_term:
            continue
        definition = sample["definition"]

        if not nlp(definition)._.is_profane:
            samples.append(sample)

    return samples


# -*- coding: utf-8 -*-
# Python 3.6
# BacteriaExp.py
"""
Author: Kebur Fantahun

Created: Thu Aug 16 10:10:24 2018
Modified: Thu Aug 16 10:10:24 2018

Description: Bacterial lab example
"""
import numpy as np, matplotlib.pyplot as plt

# Same as the comment above but for non-(csv,txt,...,etc.) formats
bacteria_data_high = np.load("g149novickA.npy")
bacteria_data_low = np.load("g149novickB.npy")

# equations
def VPop(t, tau):
    return 1 - np.exp(-t/tau)
def WPop(t, A, tau):
    return A * (np.exp(-t/tau) -1 + (t/tau))

# setup arrays for olotting
t = np.arange(0,7,3e-3)
timeH=[]
timeL= []
proteinActH=[]
proteinActL = []

# fill high inducer data arrays
for i in range(len(bacteria_data_high)):
    proteinActH.append(bacteria_data_high[i][1])
    timeH.append(bacteria_data_high[i][0])

# fill low inducer data arrays
for j in range(len(bacteria_data_low)):
    proteinActL.append(bacteria_data_low[j][1])
    timeL.append(bacteria_data_low[j][0])

# clear plot for re-running program
plt.gcf().clear()

# Experimental Model with W(t) plot
# plt.plot(t, WPop(t, 1e3, 1e3), 'bo', label=r'$A = 1  ,\tau = 1e2$')
plt.semilogy(t, WPop(t, 0.35, 1), 'r--', label=r'$A = 0.35  ,\tau = 1$')
plt.semilogy(WPop(t, 6e2, 1)-3 + 2.3, t-.1,  'c-' , label=r'$A = 6e2  ,\tau = 1$')

# How can I improve these curves? ...

# limits
plt.xlim(0.1,10)
plt.ylim(0.01,10)

# Data plot; high inducer
plt.semilogy(timeH, proteinActH, 'yx', label='High inducer')

# Data plot; low inducer
plt.semilogy(timeL, proteinActL, 'k+', label='Low inducer')

plt.legend()

ax = plt.gca()
ax.set_title("Bacterial culture production", size=16, weight='bold')

ax.set_xlabel('Time(hours)')
ax.set_ylabel('Fraction of maximum beta-galactosidase activity')

"""

Experimental Model with V(t)

plt.plot(t, 1e1*VPop(t, 5e1), 'r--', label=r'$\tau =  5e1$')
plt.plot(t, 4e2*VPop(t, 1e5), 'c-' , label=r'$\tau =  1e5$')
plt.plot(t, 1e1*VPop(t, 1e10), 'b-.', label=r'$\tau = 90$')
plt.plot(t, 1e1*VPop(t, 1e15), 'm:', label=r'$\tau = 180$')

Experimental model with W(t)

plt.plot(t, WPop(t, 1 , 1), 'r--', label=r'$A = 1  ,\tau = 1$')
plt.plot(t, WPop(t, 1, -1), 'c-', label=r'$A = 1  , \tau = -1$')
plt.plot(t, WPop(t, -1, 1), 'b-.',label=r'$A = -1,\tau = 1$')
plt.plot(t, WPop(t, -1, -1), 'r:',label=r'$A = -1,\tau = -1$')

...

Array(W#) method for A/ tau values:

W1 = [1, 1]
W2 = [1, -1]
W3 = [-1, 1]
W4 = [-1, -1]

plt.plot(t, WPop(t, W1[0] , W1[1]),'r--', t, WPop(t, W2[0] , W2[1]), 'c-',\
         t, WPop(t, W3[0] , W3[1]),'b-.', t, WPop(t, W4[0] , W4[1]), 'r:' )

"""


# coding: utf-8

class Solution(object):

    def solution(self, s):
        exists = {}
        s_len = 0
        left = 0
        for i in range(len(s)):
            if s[i] in exists and exists[s[i]] >= left:
                 left = exists[s[i]] + 1

            exists[s[i]] = i
            s_len = max(s_len, i - left + 1)

        return s_len


from jinja2 import Environment, FileSystemLoader, select_autoescape
import yaml

with open('src/content.yaml') as f:
    content = yaml.load(f, Loader=yaml.FullLoader)

env = Environment(
    loader=FileSystemLoader('src/templates/'),
    autoescape=select_autoescape(['html', 'xml'])
)
template = env.get_template('index.html')
rendered = template.render(content)

with open('_build/index.html', 'w') as f:
    f.write(rendered)


from typing import List
import torch
from torch import nn
import numpy as np

def make_pad_mask(lengths, xs=None, length_dim=-1, maxlen=None):
    """Make mask tensor containing indices of padded part.

    Args:
        lengths (LongTensor or List): Batch of lengths (B,).
        xs (Tensor, optional): The reference tensor.
            If set, masks will be the same shape as this tensor.
        length_dim (int, optional): Dimension indicator of the above tensor.
            See the example.

    Returns:
        Tensor: Mask tensor containing indices of padded part.
                dtype=torch.uint8 in PyTorch 1.2-
                dtype=torch.bool in PyTorch 1.2+ (including 1.2)

    Examples:
        With only lengths.

        >>> lengths = [5, 3, 2]
        >>> make_pad_mask(lengths)
        masks = [[0, 0, 0, 0 ,0],
                 [0, 0, 0, 1, 1],
                 [0, 0, 1, 1, 1]]

        With the reference tensor.

        >>> xs = torch.zeros((3, 2, 4))
        >>> make_pad_mask(lengths, xs)
        tensor([[[0, 0, 0, 0],
                 [0, 0, 0, 0]],
                [[0, 0, 0, 1],
                 [0, 0, 0, 1]],
                [[0, 0, 1, 1],
                 [0, 0, 1, 1]]], dtype=torch.uint8)
        >>> xs = torch.zeros((3, 2, 6))
        >>> make_pad_mask(lengths, xs)
        tensor([[[0, 0, 0, 0, 0, 1],
                 [0, 0, 0, 0, 0, 1]],
                [[0, 0, 0, 1, 1, 1],
                 [0, 0, 0, 1, 1, 1]],
                [[0, 0, 1, 1, 1, 1],
                 [0, 0, 1, 1, 1, 1]]], dtype=torch.uint8)

        With the reference tensor and dimension indicator.

        >>> xs = torch.zeros((3, 6, 6))
        >>> make_pad_mask(lengths, xs, 1)
        tensor([[[0, 0, 0, 0, 0, 0],
                 [0, 0, 0, 0, 0, 0],
                 [0, 0, 0, 0, 0, 0],
                 [0, 0, 0, 0, 0, 0],
                 [0, 0, 0, 0, 0, 0],
                 [1, 1, 1, 1, 1, 1]],
                [[0, 0, 0, 0, 0, 0],
                 [0, 0, 0, 0, 0, 0],
                 [0, 0, 0, 0, 0, 0],
                 [1, 1, 1, 1, 1, 1],
                 [1, 1, 1, 1, 1, 1],
                 [1, 1, 1, 1, 1, 1]],
                [[0, 0, 0, 0, 0, 0],
                 [0, 0, 0, 0, 0, 0],
                 [1, 1, 1, 1, 1, 1],
                 [1, 1, 1, 1, 1, 1],
                 [1, 1, 1, 1, 1, 1],
                 [1, 1, 1, 1, 1, 1]]], dtype=torch.uint8)
        >>> make_pad_mask(lengths, xs, 2)
        tensor([[[0, 0, 0, 0, 0, 1],
                 [0, 0, 0, 0, 0, 1],
                 [0, 0, 0, 0, 0, 1],
                 [0, 0, 0, 0, 0, 1],
                 [0, 0, 0, 0, 0, 1],
                 [0, 0, 0, 0, 0, 1]],
                [[0, 0, 0, 1, 1, 1],
                 [0, 0, 0, 1, 1, 1],
                 [0, 0, 0, 1, 1, 1],
                 [0, 0, 0, 1, 1, 1],
                 [0, 0, 0, 1, 1, 1],
                 [0, 0, 0, 1, 1, 1]],
                [[0, 0, 1, 1, 1, 1],
                 [0, 0, 1, 1, 1, 1],
                 [0, 0, 1, 1, 1, 1],
                 [0, 0, 1, 1, 1, 1],
                 [0, 0, 1, 1, 1, 1],
                 [0, 0, 1, 1, 1, 1]]], dtype=torch.uint8)

    """
    if length_dim == 0:
        raise ValueError(f"length_dim cannot be 0: {length_dim}")

    if not isinstance(lengths, list):
        lengths = lengths.long().tolist()

    bs = int(len(lengths))
    if maxlen is None:
        if xs is None:
            maxlen = int(max(lengths))
        else:
            maxlen = xs.size(length_dim)
    else:
        assert xs is None
        assert maxlen >= int(max(lengths))

    seq_range = torch.arange(0, maxlen, dtype=torch.int64)
    seq_range_expand = seq_range.unsqueeze(0).expand(bs, maxlen)
    seq_length_expand = seq_range_expand.new(lengths).unsqueeze(-1)
    mask = seq_range_expand >= seq_length_expand

    if xs is not None:
        assert xs.size(0) == bs, (xs.size(0), bs)

        if length_dim < 0:
            length_dim = xs.dim() + length_dim
        # ind = (:, None, ..., None, :, , None, ..., None)
        ind = tuple(
            slice(None) if i in (0, length_dim) else None for i in range(xs.dim())
        )
        mask = mask[ind].expand_as(xs).to(xs.device)
    return mask

class BLIP_VQA(nn.Module):
    def __init__(self,
                 med_config='configs/med_config.json',
                 image_size=480,
                 cased=True,
                 vit='base',
                 vit_grad_ckpt=False,
                 vit_ckpt_layer=0,
                 multitask_qcate=True,
                 ):
        """
        Args:
            med_config (str): path for the mixture of encoder-decoder model's configuration file
            image_size (int): input image size
            vit (str): model size of vision transformer
        """
        from models.med import BertConfig, BertModel, BertLMHeadModel
        from models.blip import create_vit, init_tokenizer

        super().__init__()

        self.visual_encoder, vision_width = create_vit(vit, image_size, vit_grad_ckpt, vit_ckpt_layer,
                                                       drop_path_rate=0.1)
        self.tokenizer = init_tokenizer(cased)

        encoder_config = BertConfig.from_json_file(med_config)
        encoder_config.encoder_width = vision_width
        self.text_encoder = BertModel(config=encoder_config, add_pooling_layer=True)

        decoder_config = BertConfig.from_json_file(med_config)
        self.text_decoder = BertLMHeadModel(config=decoder_config)

        self.num_heads = encoder_config.num_attention_heads
        self.num_patches = self.visual_encoder.patch_embed.num_patches + 1
        self.multitask_qcate = multitask_qcate
        if multitask_qcate:
            # self.retr_ffn = nn.Linear(self.num_patches, 1)
            self.multitask_ffn = nn.Linear(encoder_config.hidden_size, 6)

    def encode_images(self, images: torch.Tensor, n_facts: List[int]):
        """
        :param images: (batch, n_facts, channel, H, W)
        :param n_facts: Number of image facts in each sample in the batch
        :return:
            - image_embeds: (batch, n_facts * seq_len, embed_size)
            - lengths: Valid lengths (dim 1) of image_embeds
        """
        batch_size, max_n_facts, C, H, W = images.shape
        images = images.view(-1, C, H, W)
        image_embeds = self.visual_encoder(images)  # (batch * n_facts, seq_len, embed_size)

        n_facts = [nf * image_embeds.size(1) for nf in n_facts]
        image_embeds = image_embeds.view(batch_size, max_n_facts * image_embeds.size(1), image_embeds.size(2))
        return image_embeds, n_facts

    def enable_med(self):
        tie_encoder_decoder_weights(self.text_encoder, self.text_decoder.bert, '', '/attention')

    def forward(
            self,
            image: torch.Tensor,
            captions: List[List[str]],
            question: List[str],
            answer: List[str],
            n_img_facts: List[int],
            train=True,
    ):
        """
        :param image: (batch, n_img_facts, channel, H, W)
        :param captions: Batch of list of captions
        :param question: Batch of questions
        :param answer: Batch of answers
        :param n_img_facts: Batch of number of image facts
        :param train: train or inference
        """

        image_embeds, lengths = self.encode_images(image, n_img_facts)
        image_atts = ~make_pad_mask(lengths, image_embeds[:, :, 0], 1).to(image.device)

        question = self.tokenizer(question, padding='longest', return_tensors="pt").to(image.device)
        question.input_ids[:, 0] = self.tokenizer.enc_token_id

        # concatenate captions and tokenize them
        captions = [f' {self.tokenizer.sep_token} '.join(cap) for cap in captions]
        captions = self.tokenizer(captions, padding='longest', return_tensors="pt").to(image.device)
        # mask the first token since we already have a sep_token_id set to the last token of question
        captions.input_ids[:, 0] = self.tokenizer.pad_token_id
        captions.attention_mask[:, 0] = 0

        # image-grounded text encoder
        input_ids = torch.cat([question.input_ids, captions.input_ids], dim=-1)
        attention_mask = torch.cat([question.attention_mask, captions.attention_mask], dim=-1)
        input_ids = input_ids[:, :512]
        attention_mask = attention_mask[:, :512]

        cross_attention_weight = torch.ones_like(input_ids, dtype=torch.float)
        cross_attention_weight[torch.as_tensor(n_img_facts, dtype=torch.long) == 0] = 0.0
        question_output = self.text_encoder(
            input_ids,
            attention_mask=attention_mask,
            encoder_hidden_states=image_embeds,
            encoder_attention_mask=image_atts,
            cross_attention_weight=cross_attention_weight,
            # output_attentions=True,
            return_dict=True,
        )

        # (batch, num_heads, question_len, image_embeds_len)
        multimodal_cross_atts = None
        if train:
            if self.multitask_qcate:
                # multimodal_cross_atts = question_output.cross_attentions[-1]  # last layer's cross attention
                # atts = torch.sum(multimodal_cross_atts, dim=2)  # (batch, num_heads, image_embeds_len)
                # atts = torch.sum(atts, dim=1)  # (batch, image_embeds_len)

                # # (batch, n_facts, num_patches)
                # atts = atts.view(atts.shape[0], -1, self.num_patches)
                # retr = self.retr_ffn(atts).squeeze(dim=-1)  # (batch, n_facts)

                mt_res = self.multitask_ffn(question_output.pooler_output)
            else:
                mt_res = None

            '''
            n: number of answers for each question
            weights: weight for each answer
            '''
            answer = self.tokenizer(answer, padding='longest', return_tensors="pt").to(image.device)
            answer.input_ids[:, 0] = self.tokenizer.bos_token_id
            answer_targets = answer.input_ids.masked_fill(answer.input_ids == self.tokenizer.pad_token_id, -100)

            answer_output = self.text_decoder(answer.input_ids,
                                              attention_mask=answer.attention_mask,
                                              encoder_hidden_states=question_output.last_hidden_state,
                                              encoder_attention_mask=attention_mask,
                                              labels=answer_targets,
                                              return_dict=True,
                                              reduction='none',
                                              )

            loss = answer_output.loss
            loss = loss.sum() / image.size(0)
            return loss, mt_res, multimodal_cross_atts
        else:
            num_beams = 10
            question_states = question_output.last_hidden_state.repeat_interleave(num_beams, dim=0)
            question_atts = attention_mask.repeat_interleave(num_beams, dim=0).to(question_states.device)

            bos_ids = torch.full((image.size(0), 1), fill_value=self.tokenizer.bos_token_id, device=image.device)

            outputs = self.text_decoder.generate(
                input_ids=bos_ids,
                max_length=50,
                min_length=1,
                num_beams=num_beams,
                eos_token_id=self.tokenizer.sep_token_id,
                pad_token_id=self.tokenizer.pad_token_id,
                encoder_hidden_states=question_states,
                encoder_attention_mask=question_atts,
            )

            answers = []
            for output in outputs:
                answer = self.tokenizer.decode(output, skip_special_tokens=True)
                answers.append(answer)
            return answers

def blip_vqa(pretrained='', **kwargs):
    from models.blip import load_checkpoint

    model = BLIP_VQA(**kwargs)
    if pretrained:
        model, msg = load_checkpoint(model, pretrained)
    return model

def tile(x, dim, n_tile):
    init_dim = x.size(dim)
    repeat_idx = [1] * x.dim()
    repeat_idx[dim] = n_tile
    x = x.repeat(*(repeat_idx))
    order_index = torch.LongTensor(np.concatenate([init_dim * np.arange(n_tile) + i for i in range(init_dim)]))
    return torch.index_select(x, dim, order_index.to(x.device))

def tie_encoder_decoder_weights(encoder: nn.Module, decoder: nn.Module, base_model_prefix: str, skip_key: str):
    """
    https://github.com/salesforce/BLIP/blob/main/models/blip_pretrain.py
    """
    uninitialized_encoder_weights: List[str] = []
    if decoder.__class__ != encoder.__class__:
        print(
            f"{decoder.__class__} and {encoder.__class__} are not equal. "
            f"In this case make sure that all encoder weights are correctly initialized."
        )

    def tie_encoder_to_decoder_recursively(
            decoder_pointer: nn.Module,
            encoder_pointer: nn.Module,
            module_name: str,
            uninitialized_encoder_weights: List[str],
            skip_key: str,
            depth=0,
    ):
        assert isinstance(decoder_pointer, nn.Module) and isinstance(
            encoder_pointer, nn.Module
        ), f"{decoder_pointer} and {encoder_pointer} have to be of type torch.nn.Module"
        if hasattr(decoder_pointer, "weight") and skip_key not in module_name:
            assert hasattr(encoder_pointer, "weight")
            encoder_pointer.weight = decoder_pointer.weight
            if hasattr(decoder_pointer, "bias"):
                assert hasattr(encoder_pointer, "bias")
                encoder_pointer.bias = decoder_pointer.bias
            print(module_name + ' is tied')
            return

        encoder_modules = encoder_pointer._modules
        decoder_modules = decoder_pointer._modules
        if len(decoder_modules) > 0:
            assert (
                    len(encoder_modules) > 0
            ), f"Encoder module {encoder_pointer} does not match decoder module {decoder_pointer}"

            all_encoder_weights = set([module_name + "/" + sub_name for sub_name in encoder_modules.keys()])
            encoder_layer_pos = 0
            for name, module in decoder_modules.items():
                if name.isdigit():
                    encoder_name = str(int(name) + encoder_layer_pos)
                    decoder_name = name
                    if not isinstance(decoder_modules[decoder_name], type(encoder_modules[encoder_name])) and len(
                            encoder_modules
                    ) != len(decoder_modules):
                        # this can happen if the name corresponds to the position in a list module list of layers
                        # in this case the decoder has added a cross-attention that the encoder does not have
                        # thus skip this step and subtract one layer pos from encoder
                        encoder_layer_pos -= 1
                        continue
                elif name not in encoder_modules:
                    continue
                elif depth > 500:
                    raise ValueError(
                        "Max depth of recursive function `tie_encoder_to_decoder` reached. "
                        "It seems that there is a circular dependency between two or more `nn.Modules` of your model."
                    )
                else:
                    decoder_name = encoder_name = name
                tie_encoder_to_decoder_recursively(
                    decoder_modules[decoder_name],
                    encoder_modules[encoder_name],
                    module_name + "/" + name,
                    uninitialized_encoder_weights,
                    skip_key,
                    depth=depth + 1,
                )
                all_encoder_weights.remove(module_name + "/" + encoder_name)

            uninitialized_encoder_weights += list(all_encoder_weights)

    # tie weights recursively
    tie_encoder_to_decoder_recursively(decoder, encoder, base_model_prefix, uninitialized_encoder_weights, skip_key)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Jul 29 08:21:14 2019

@author: Sergio Ribeiro
Description: Problem Definitions: Input Dictionary
"""
from dolfin import *
import numpy as np
import math

# Sigmoid Function S(x) = (dtMax-dtMin)/(1+exp(Inclination*(-t+t0)))+dtMin
def dynamicTimestep(t,dtMax,dtMin,t0):
    Inclination = 3
    return (dtMax-dtMin)/(1+math.exp(Inclination*(-t+t0)))+dtMin

def autoTimestep(no_iterations,dt,inputs,limitIterations=4,increment=2):
    # Check if
    if no_iterations < limitIterations:
        dt = min(increment*dt,inputs.dtMax)
    elif no_iterations > limitIterations + 2:
        dt = max((1/increment)*dt,inputs.dtMin)
    else:
        dt = dt

    return dt

def dynamicSaveDt(dt):
    return 5*dt

class Inputs():
    def __init__(self):

        #%%############ Problem Geometry   ##############################
        ## Mesh File
        self.meshFile = 'WellSimulatorRZ'#'MacroParallelPlates'#'WellSimulator'#
        self.meshElements = 8813
        # Geometric Values
        self.Zmin = 6
        self.ZVelProf = 6.5
        self.ZFL = 7
        self.Zmax = 8
        self.RIn = 0.1
        self.stepR = 0.005
        self.RIn = 0.1
        self.ROut = 0.22
        self.HFluidLoss = .1
        ## Mesh Elements
        # Velocity
        self.velocityElementfamily = 'Lagrange'
        self.velocityElementOrder = 2
        # Pressure
        self.pressureElementfamily = 'Lagrange'
        self.pressureElementOrder = 1
        # Advected Scalar Field
        self.scalarFieldElementfamily = 'Lagrange'
        self.scalarFieldElementOrder = 1

        #%%############ Case Definition    ##############################
        self.caseId = 'TransWellSimulatorRZ_8Cyl_MeshTest_MF10' ## If name already exists in folder ./PostProcessing/Cases,
                    ## old data will be overwritten.

        # Output Variables
        self.ParaViewFilenames = []; self.ParaViewTitles = []
        self.ParaViewFilenames.append("velocity"); self.ParaViewTitles.append('Velocity (m/s)')
        self.ParaViewFilenames.append("pressure"); self.ParaViewTitles.append('Pressure (Pa)')
        self.ParaViewFilenames.append("concentration"); self.ParaViewTitles.append('Mass Fraction (Fluid Tags)')
        self.ParaViewFilenames.append("density"); self.ParaViewTitles.append('Density (kg/m³)')
        self.ParaViewFilenames.append("viscosity"); self.ParaViewTitles.append('Viscosity (Pa.s)')
        self.outputFlowrate = './PostProcessing/Cases/'+self.caseId+'/flowrateInput.csv'
        self.outputPressure = './PostProcessing/Cases/'+self.caseId+'/pressureOutput.csv'
        self.outputTOC = './PostProcessing/Cases/'+self.caseId+'/TOC.csv'

        #%%############ Time Parameters #################################
        # Start Time
        self.t0 = 0 # s

        # Simulation Time
        self.tEnd = 200 # s

        # MeshTest Radius List
        self.plotRadiusList = np.arange(self.RIn,self.ROut+self.stepR,self.stepR).tolist()
        # Plot Time List
        self.plotTimeList = [1, 10, 50, 100, 150, self.tEnd]# 2500, 4500, 7000, 8250, 9000, 10500, 12000,
        self.plotDepthList = [6, 6.4, 6.8, 7.2, 7.6, 8]
        # self.plotTimeList = [1, self.tEnd]
        self.fieldnamesFlow = ['Time(s)','outletFlowRate(Kg/s)']
        self.dZPlot = 0.01
        self.fieldnamesPre = ['Time(s)','rhoInlet(Kg/m3)']

        # Variable time step
        self.dtMin = 1e-4   # s
        self.dtMax = 2e1    # s
        self.tChange = 0    # point in time of sigmoid inflection occurs (s)
        self.dt = 1e-3      # s
#        self.dt = dynamicTimestep(self.t0,self.dtMax,self.gging Options   ###############################

        # Result Saving time step
        self.savedt = 100  # s

        #%%############ Gravitationa Field ##############################
        # Gravity Acceleration (m/s²) on axis X
        self.g = 9.81

        #%%############ Fluids' Properties ##############################
        # Tags
        self.Fluid0 = 0 # Cement
        self.Fluid1 = 1 - self.Fluid0 # Water
        self.CInitialMixture = 0.5      # Mass fraction of Fluid0. Fluid1 = 1-Flui0

        # Diffusity of Between species (m²/s)
        self.D = 1e-3

        # Rheology
        # Newtonian Viscosity
        self.mu_water = 0.01    # Pa.s
        self.mu_cem = 0.3       # Pa.s
        self.mu_values = [self.mu_cem , self.mu_water]

        ## Rheology - Modified SMD (Souza Mendes e Dutra (2004)) + Cure(tauY(t))
        # Input Variables
        self.tau0 = 19.019         # Dinamic Yield Stress
        self.eta0 = 1.3e4            # Viscosity Value for Low shear rates
        self.etaInf = 0.1*self.eta0  # Equilibrium Viscosity(Newtonian Plato: Lowgh shear rates)
        # self.etaInf = self.mu_cem       # Equilibrium Viscosity(Newtonian Plato: Lowgh shear rates)
        self.K = 1.43              # Consistency Index
        self.n = 0.572             # Power-law Index
        self.ts = 6000             # Caracteristic viscosity buildup time
        self.eps = 1e-10
        self.tauY_t = Expression('tau0*exp(t/ts)',degree=1,\
                                tau0 = self.tau0, \
                                ts = self.ts, t = 0)
        # Modified SMD Model Expression
        self.rheologicalModel = Expression('(1 - exp(-eta0*(gammaDot)/tauY_t))* \
                                            (tauY_t/(gammaDot+eps) + K*(pow(gammaDot+eps,nPow)/ \
                                            (gammaDot+eps))) \
                                            + etaInf', degree=2, \
                                            etaInf=self.etaInf, eta0=self.eta0, \
                                            K = self.K, nPow = self.n, ts = self.ts, \
                                            eps = self.eps, gammaDot = Constant(1/self.eps),\
                                            tauY_t = Constant(self.eps))

        # Density (kg/m³)
        # Experimental Values
        self.rho_water = 1000           # kg/m³
        self.rho_bulk0 = 1845.48        # kg/m³ = 15.4ppg
        self.rho_bulkInf = 1200         # kg/m³
        self.rho_cem0 = (self.rho_bulk0 - (1-self.CInitialMixture)*self.rho_water)/(self.CInitialMixture)
        self.rho_cemInf = (self.rho_bulkInf - (1-self.CInitialMixture)*self.rho_water)/(self.CInitialMixture)

        # Initial Density Values per Component
        self.rho_values = [self.rho_cem0 , self.rho_water] # kg/m³

        # Shrinkage Equation: rhoMax - ((rhoMax-rhoMin)/(1+math.exp(Inclination*(-t+t0)))+rhoMin) +rhoMin
        # if Inclination is zero, shrinkage is neglected
        self.shrinkage_inclination = 0.0004         # kg/m³ / s
        self.shrinkage_rhoMax = self.rho_cem0     # kg/m³
        self.shrinkage_rhoMin = self.rho_cemInf    # kg/m³
        self.shrinkage_t0 = 1.2*self.ts             # s

        # Shrinkage Model Expression
        self.shrinkageModel = Expression('(rhoMax-((rhoMax-rhoMin)/  \
                                        (1 + exp(Inclination*(-t+t0))) \
                                        +rhoMin) + rhoMin)', degree=2, \
                                        rhoMax = self.shrinkage_rhoMax, \
                                        rhoMin = self.shrinkage_rhoMin, \
                                        Inclination = self.shrinkage_inclination, \
                                        t=self.t0, t0 = self.shrinkage_t0)
        # self.shrinkageModel = Expression('(rhoMax-((rhoMax-rhoMin)/  \
        #                                 (1 + exp(Inclination*(-t+t0))) \
        #                                 +rhoMin) + rhoMin)*cFrac + rho_water*(1-cFrac)', degree=2, \
        #                                 cFrac = Constant(self.eps), \
        #                                 rhoMax = self.rho_values[self.Fluid0], \
        #                                 rhoMin = self.shrinkage_rhoMin, \
        #                                 Inclination = self.shrinkage_inclination, \
        #                                 t=self.t0, t0 = self.shrinkage_t0, \
        #                                 rho_water = self.rho_values[self.Fluid1])

        #%%############ Boundary Conditions

        ## No slip Boundaries
        self.noSlipBCs = []
        # self.noSlipBCs.append('TopWall') # Both
        self.noSlipBCs.append('InnerPipe') # WelSimulator only
        self.noSlipBCs.append('OuterWall') # WelSimulator only
        self.noSlipBCs.append('BottomWall')  # Both
        #noSlipBoundaries.append('InnerWalls')

        ## Advected Scalar Field Inputs
        self.CInitialMixture = 0.5      # Mass fraction of Fluid0. Fluid1 = 1-Flui0
        self.scalarFieldBCs = {}
        self.COutlet = 1
        self.scalarFieldBCs.update({'Inlet' : Constant(self.CInitialMixture)}) # CMix
        # self.scalarFieldBCs.update({'Outlet' : Constant(self.CInitialMixture)}) # CMix: REGULAR FLOW
        self.scalarFieldBCs.update({'Outlet': Constant(self.COutlet)}) # C1: FILTRATION

        ## Velocity Inputs
        t=self.dtMin
        self.AOut = 2*pi*self.ROut*self.HFluidLoss
        self.AIn = pi*(self.ROut**2 - self.RIn**2)
        self.rhoOut = self.rho_values[self.Fluid1] # Water Only
        self.muOut = self.mu_values[self.Fluid1] # Water Only
        self.rho_mix0 = (self.rho_values[0]*self.CInitialMixture + self.rho_values[1]*(1-self.CInitialMixture))
        self.velocityBCs = {}
        # self.VrOutlet = '0.00043 + 0*t*A*rho'
        # self.VrOutlet = 't <= 100 ? (1/(rho*A))*0.00163 : (1/(rho*A))*0.061/((pow(t,0.78)))' # BaseCase 0
        self.VzOutlet = Expression('VzOutlet',VzOutlet=0.0,degree=1)
        self.VrInlet = Expression('VrInlet',VrInlet=0.0,degree=2)
        # self.VrOutlet = Expression('(1/(rho*A))*0.000163 + 0*t',\
        #                     A=self.AOut,rho=self.rhoOut,t=t,degree=2) # Constant Flow Rate
        self.VrOutlet = Expression('t <= 100 ? (1/(rho*A))*0.00163 : (1/(rho*A))*(a/(pow(t,b)) - (1/(d*c))*exp((t-ts)/(e*c)))',\
                                    A=self.AOut,rho=self.rhoOut,t=t,degree=2, \
                                    a= 0.032,b=0.67 ,c=1300, d=8, e=3 ,ts=17000) # Best Experimental Fit
        # self.VzInlet = Expression('t <= 100 ? (1/(rho*A))*0.00163 : (1/(rho*A))*(a/(pow(t,b)) - (1/(d*c))*exp((t-ts)/(e*c)))',\
        #                             A=self.AIn,rho=self.rho_mix0,t=t,degree=2, \
        #                             a= 0.032,b=0.67 ,c=1300, d=8, e=3 ,ts=17000) # Best Experimental Fit
        # self.VrOutlet = Expression('t <= 100 ? (1/(rho*A))*0.000163 : (1/(rho*A))*0.0055 /((pow(t,0.78)))',\
        #                            A=self.AOut,rho=self.rhoOut,t=t,degree=2) # All BaseCases
        # self.VrOutlet = 't <= 100 ? (1/(rho*A))*0.0163 : (1/(rho*A))*0.55 /((pow(t,0.78)))'
        # self.VrOutlet = 't <= 100 ? (1/(rho*A))*0.000163 : (1/(rho*A))*0.0055 /((pow(t,0.78)))'#'2*exp(1-(t/200))/300'#'2*exp(1-(t/200))/300'#
        # self.velocityBCs.update({'Inlet' : {2: self.VrInlet}}) # m/s
        InletVelBC = {0: self.VrInlet}
        # InletVelBC.update({1: self.VzInlet})
        self.velocityBCs.update({'Inlet' : InletVelBC}) # m/s
        # self.velocityBCs.update({'Outlet' : {1: self.VzOutlet}}) # m/s
        OutletVelBC = {0: self.VrOutlet}
        OutletVelBC.update({1: self.VzOutlet})
        self.velocityBCs.update({'Outlet' : OutletVelBC}) # m/s

        ## Pressure Inputs
        self.pressureBCs = {}
        # Inlet Pressure
        self.pInlet = self.rho_mix0*self.Zmin*self.g #0.3164557 #self.rho_values[0]*2*self.g
        self.pressureBCs.update({'Inlet' : self.pInlet}) # Pa
        # Outlet Pressure
        # self.pOutlet = 0.60*(self.rho_mix0*self.g*self.ZFL)
        # self.pressureBCs.update({'Outlet' : self.pOutlet}) # Pa

        #%%############ Solver parameters ###############################
        # Absolute Tolerance
        self.absTol = 1e-12

        # Relative Tolerance
        self.relTol = 1e-13

        # Maximum Iterations
        self.maxIter = 15

        # Linear Solver
        self.linearSolver = 'mumps'

        # Relaxation Factors
        self.alpha = 0.9
        self.alphaC = 0.9

        #%% Possible Solvers
        # Solver method  |  Description
        # -----------------------------------------------------------
        # bicgstab       |  Biconjugate gradient stabilized method
        # cg             |  Conjugate gradient method
        # default        |  default linear solver
        # gmres          |  Generalized minimal residual method
        # minres         |  Minimal residual method
        # mumps          |  MUMPS (MUltifrontal Massively Parallel Sparse direct Solver)
        # petsc          |  PETSc built in LU solver
        # richardson     |  Richardson method
        # superlu        |  SuperLU
        # tfqmr          |  Transpose-free quasi-minimal residual method
        # umfpack        |  UMFPACK (Unsymmetric MultiFrontal sparse LU factorization)
        #
        # LU method  |  Description
        # --------------------------------------------------------------------------
        # default    |  default LU solver
        # mumps      |  MUMPS (MUltifrontal Massively Parallel Sparse direct Solver)
        # petsc      |  PETSc built in LU solver
        # superlu    |  SuperLU
        # umfpack    |  UMFPACK (Unsymmetric MultiFrontal sparse LU factorization)
        #
        # Krylov method  |  Description
        # --------------------------------------------------------------
        # bicgstab       |  Biconjugate gradient stabilized method
        # cg             |  Conjugate gradient method
        # default        |  default Krylov method
        # gmres          |  Generalized minimal residual method
        # minres         |  Minimal residual method
        # richardson     |  Richardson method
        # tfqmr          |  Transpose-free quasi-minimal residual method
        #
        #


import torch
import torch.nn as nn
from torch.utils.tensorboard import SummaryWriter
from .utils import count_active_neighbors
from torch.nn.functional import interpolate

class LoggingController:
    img_epochs = 10

    def __init__(self, model, writer: SummaryWriter):
        self.model = model
        self.writer = writer
        self._all_logging_layers = None

    def log_densities(self, epoch, reset=False, disable=False, prefix=""):
        for layer in self.get_all_logging_layers().get(DensityLogger, []):
            self.writer.add_scalar(f"{prefix}densities/{layer.name}", layer.get(), epoch)
            if reset:
                layer.reset()
            if disable:
                layer.disable_logging()

    def _log_pixel_masks_impl(self, channels, name, epoch, individual_channels=False, prefix=""):
        def to_img(x):
            return (x * 255).to(torch.uint8)

        self.writer.add_image(f"{prefix}pixel-mask-acc/{name}/", to_img(torch.mean(channels, dim=0, keepdim=True)), epoch)

        if individual_channels:
            for ch_idx in range(len(channels)):
                self.writer.add_image(f"{prefix}pixel-mask-channels/{name}/{ch_idx}/", to_img(channels[None, ch_idx]), epoch)

    def log_pixel_masks(self, epoch, reset=False, disable=False):
        for layer in self.get_all_logging_layers().get(PixelMaskLogger, []):
            self._log_pixel_masks_impl(layer.get(), layer.name, epoch)

            if reset:
                layer.reset()
            if disable:
                layer.disable_logging()

    def log_diff_in_histogram(self, epoch, reset=False, disable=False):
        total_histogram = []
        for layer in self.get_all_logging_layers().get(DifferentialPixelMaskLogger, []):
            histrogram = layer.get_histogram(tile_size=(6, 6))

            bins = range(0, 65)

            if histrogram is not None and len(histrogram) > 0:
                self.writer.add_histogram(f"tile-density/{layer.name}", histrogram, epoch, bins=bins)
                histogram_non_empty = histrogram[histrogram > 0]
                if len(histogram_non_empty > 0):
                    self.writer.add_histogram(f"tile-density-non-empty/{layer.name}", histogram_non_empty, epoch, bins=bins)
                    total_histogram.extend(list(histrogram))

            if reset:
                layer.reset()
            if disable:
                layer.disable_logging()

        if len(total_histogram) > 0:
            t_total_histogram = torch.tensor(total_histogram)
            self.writer.add_histogram(f"tile-density-total/all", t_total_histogram, epoch, bins=bins)
            t_total_histogram_non_empty = t_total_histogram[t_total_histogram > 0]
            if len(t_total_histogram_non_empty) > 0:
                self.writer.add_histogram(f"tile-density-total/non-empty", t_total_histogram_non_empty, epoch, bins=bins)

    def image_logger(self, epoch, reset, disable, path, cls):
        def to_img(x):
            return (x * 255).to(torch.uint8)

        for layer in self.get_all_logging_layers().get(cls, []):
            if epoch % LoggingController.img_epochs == (LoggingController.img_epochs - 1):
                val = layer.get()
                if val is not None:
                    self.writer.add_image(f"{path}/{layer.name}/", to_img(torch.mean(val[None], dim=0, keepdim=True)), epoch)
            if reset:
                layer.reset()
            if disable:
                layer.disable_logging()

    def log_prev_input(self, epoch, reset=False, disable=False):
        self.image_logger(epoch, reset, disable, "prev_input", PrevInputLogger)

    def log_input(self, epoch, reset=False, disable=False):
        self.image_logger(epoch, reset, disable, "input", InputLogger)

    def log_output(self, epoch, reset=False, disable=False):
        self.image_logger(epoch, reset, disable, "output", OutputLogger)

    def log_computations(self, epoch, reset=False, disable=False):
        def to_img(x):
            return (x * 255).to(torch.uint8)

        all_masks = []
        max_res = [0,0]

        logged_total_computations = False
        for layer in self.get_all_logging_layers().get(ComputationsLogger, []):
            if not logged_total_computations:
                self.writer.add_scalar(f"computations", layer.get(), epoch)
                logged_total_computations = True

            mask = layer.get_mask()
            if mask is not None and epoch % LoggingController.img_epochs == (LoggingController.img_epochs - 1):
                self.writer.add_image(f"update-mask/{layer.name}/", to_img(torch.mean(mask[None], dim=0, keepdim=True)), epoch)
                all_masks.append(mask)
                if mask.shape[0] > max_res[0] or mask.shape[1] > max_res[1]:
                    max_res = mask.shape

            if reset:
                layer.reset()
            if disable:
                layer.disable_logging()

        if len(all_masks) > 0:
            for i in range(len(all_masks)):
                all_masks[i] = interpolate(all_masks[i][None, None], size=(max_res[-2], max_res[-1]), mode="nearest")
            mean_mask = to_img(torch.mean(torch.cat(all_masks, dim=1),dim=1))
            self.writer.add_image("update-density", mean_mask, epoch)

    def log_multiplications(self, epoch, reset=False, disable=False):
        logged_global = False
        for layer in self.get_all_logging_layers().get(MultiplicationsLogger, []):
            if not logged_global:
                self.writer.add_scalar(f"multiplications/rel", layer.get_global_rel(), epoch)
                self.writer.add_scalar(f"multiplications/abs", layer.get_global_abs(), epoch)
                logged_global = True

            self.writer.add_scalar(f"mul-per-layer-rel/{layer.name}", layer.get_self_rel(), epoch)
            self.writer.add_scalar(f"mul-per-layer-abs/{layer.name}", layer.get_self_abs(), epoch)

            if reset:
                layer.reset()
            if disable:
                layer.disable_logging()

    def log_prev_input_bandwidth(self, epoch, reset=False, disable=False):
        for layer in self.get_all_logging_layers().get(PrevInBandwidthLogger, []):
            self.writer.add_scalar(f"prev_in_bandwidth", layer.get(), epoch)

            if reset:
                layer.reset()
            if disable:
                layer.disable_logging()
            break

    def log_prev_in_sizes(self, epoch, reset=False, disable=False):
        logged_global = False
        for layer in self.get_all_logging_layers().get(PrevInputSizeLogger, []):
            if not logged_global:
                self.writer.add_scalar(f"prev-in-size-global/", layer.get_global(), epoch)
                logged_global = True

            self.writer.add_scalar(f"prev-in-size/{layer.name}/", layer.get(), epoch)

            if reset:
                layer.reset()
            if disable:
                layer.disable_logging()

    def log_gradient(self, epoch, reset=False, disable=False):
        for layer in self.get_all_logging_layers().get(GradientLogger, []):
            result = layer.get()
            if result is None:
                continue
            self.writer.add_image(f"gradient", result, epoch)

            if reset:
                layer.reset()
            if disable:
                layer.disable_logging()

    def log_combined_loggers(self, epoch, reset=False, disable=False):
        logged_total_diff_density = False

        for layer in self.get_all_logging_layers().get(Loggers, []):
            outputs = layer.get()

            for log_name, log_out in outputs.items():
                if "DiffMask" in log_name and epoch % LoggingController.img_epochs == (LoggingController.img_epochs - 1):
                    self._log_pixel_masks_impl(log_out, layer.name, epoch, prefix="diff-")
                elif "Mask" in log_name and epoch % LoggingController.img_epochs == (LoggingController.img_epochs - 1):
                    self._log_pixel_masks_impl(log_out, layer.name, epoch)
                elif "DiffDensity" in log_name:
                    self.writer.add_scalar(f"diff-densities/{layer.name}", log_out, epoch)
                    if not logged_total_diff_density:
                        self.writer.add_scalar(f"diff_densities_sum", layer.logs[log_name].get_sum(), epoch)
                        logged_total_diff_density = True
                elif "Density" in log_name:
                    self.writer.add_scalar(f"densities/{layer.name}", log_out, epoch)

            if reset:
                layer.reset()
            if disable:
                layer.disable_logging()

    def get_all_logging_layers(self):
        # if list was already computed, return a copy of it to save time
        if self._all_logging_layers is not None:
            return {**self._all_logging_layers}

        layers = {}

        def recursive_get_layers(module):
            for layer in module.modules():
                cls = type(layer)

                is_logging_layer = issubclass(cls, LoggingLayer)
                if is_logging_layer and cls not in layers:
                    layers[cls] = []
                if is_logging_layer and layer in layers[cls]:
                    continue
                elif is_logging_layer:
                    layers[cls].append(layer)
                elif layer != module:
                    recursive_get_layers(layer)

        recursive_get_layers(self.model)
        self._all_logging_layers = {**layers}
        return layers

    def reset_loggers(self, enable=False, disable=False):
        for layers in self.get_all_logging_layers().values():
            for layer in layers:
                layer.reset()
                if enable:
                    layer.enable_logging()
                if disable:
                    layer.disable_logging()

    def reset_loggers_history(self):
        for layers in self.get_all_logging_layers().values():
            for layer in layers:
                layer.reset_history()

    def write_logs(self, epoch, reset=False, disable=False):
        self.log_densities(epoch, reset, disable)
        self.log_pixel_masks(epoch, reset, disable)
        self.log_computations(epoch, reset, disable)
        self.log_multiplications(epoch, reset, disable)
        self.log_prev_input_bandwidth(epoch, reset, disable)
        self.log_input(epoch, reset, disable)
        self.log_output(epoch, reset, disable)
        self.log_prev_input(epoch, reset, disable)
        self.log_prev_in_sizes(epoch, reset, disable)
        self.log_diff_in_histogram(epoch, reset, disable)
        # self.log_combined_loggers(epoch, reset, disable)

class LoggingLayer(nn.Module):
    id = 1

    def __init__(self, name="", enabled=False):
        super(LoggingLayer, self).__init__()
        self.name = name
        self.enabled = enabled
        self.added_id = False

    def enable_logging(self):
        self.enabled = True

    def disable_logging(self):
        self.enabled = False

    def get(self):
        pass

    def reset(self):
        pass

    def reset_history(self):
        pass

class DensityLogger(LoggingLayer):
    def __init__(self, threshold=0.0, **kwargs):
        super(DensityLogger, self).__init__(**kwargs)
        self.active_pixels = 0
        self.total_pixels = 0
        self.threshold = threshold

    def forward(self, x):
        if not self.added_id:
            self.name = f"{LoggingLayer.id} {self.name}"
            LoggingLayer.id += 1
            self.added_id = True

        if self.enabled:
            self.active_pixels += torch.sum(torch.abs(x) > self.threshold)
            self.total_pixels += x.numel()
        return x

    def get(self):
        return float(self.active_pixels) / float(self.total_pixels)

    def reset(self):
        self.active_pixels = 0
        self.total_pixels = 0

class DifferentialDensityLogger(LoggingLayer):
    sum_active_pixels = 0
    sum_total_pixels = 0

    def __init__(self, threshold=0.05, **kwargs):
        super(DifferentialDensityLogger, self).__init__(**kwargs)
        self.active_pixels = 0
        self.total_pixels = 0
        self.prev_frame = None
        self.threshold = threshold

    def forward(self, x):
        if not self.added_id:
            self.name = f"{LoggingLayer.id} {self.name}"
            LoggingLayer.id += 1
            self.added_id = True

        if self.enabled:
            if self.prev_frame is None:
                diff = torch.ones_like(x)
            else:
                diff = torch.abs(self.prev_frame - x) > self.threshold
            self.active_pixels += torch.sum(diff != 0.0)
            self.total_pixels += diff.numel()
            self.prev_frame = x.detach()
            DifferentialDensityLogger.sum_active_pixels += torch.sum(diff != 0.0)
            DifferentialDensityLogger.sum_total_pixels += diff.numel()
        return x

    def get(self):
        return float(self.active_pixels) / float(self.total_pixels)

    def get_sum(self):
        if DifferentialDensityLogger.sum_total_pixels == 0:
            return 0
        return float(DifferentialDensityLogger.sum_active_pixels) / float(DifferentialDensityLogger.sum_total_pixels)

    def reset(self):
        self.active_pixels = 0
        self.total_pixels = 0
        DifferentialDensityLogger.sum_active_pixels = 0
        DifferentialDensityLogger.sum_total_pixels = 0

    def reset_history(self):
        self.prev_frame = None

class PixelMaskLogger(LoggingLayer):
    def __init__(self, threshold=0.0, **kwargs):
        super(PixelMaskLogger, self).__init__(**kwargs)
        self.channels = None
        self.n_samples = 0
        self.threshold = threshold

    def reset(self):
        self.channels = None
        self.n_samples = 0

    def forward(self, x):
        if not self.added_id:
            self.name = f"{LoggingLayer.id} {self.name}"
            LoggingLayer.id += 1
            self.added_id = True

        if self.enabled:
            if self.channels is None:
                self.channels = torch.zeros_like(x[0])
            self.channels += torch.sum(torch.abs(x) > self.threshold, dim=0)
            self.n_samples += len(x)

        return x

    def get(self):
        return self.channels / self.n_samples

class DifferentialPixelMaskLogger(LoggingLayer):
    def __init__(self, threshold=0.05, **kwargs):
        super(DifferentialPixelMaskLogger, self).__init__(**kwargs)
        self.channels = None
        self.n_samples = 0
        self.prev_frame = None
        self.threshold = threshold

    def reset(self):
        self.channels = None
        self.n_samples = 0

    def reset_history(self):
        self.prev_frame = None

    def forward(self, x, use_direct=True):
        if not self.added_id:
            self.name = f"{LoggingLayer.id} {self.name}"
            LoggingLayer.id += 1
            self.added_id = True

        if self.enabled:
            if self.channels is None:
                self.channels = torch.zeros_like(x[0]).int()

            if use_direct:
                diff = x
            else:
                if self.prev_frame is None:
                    diff = torch.ones_like(x)
                else:
                    diff = torch.abs(self.prev_frame - x) > self.threshold
            self.channels += torch.sum(diff, dim=0)
            self.n_samples += len(x)
            if not use_direct:
                self.prev_frame = x.detach()

        return x

    def get(self):
        return self.channels / self.n_samples

    def get_histogram(self, tile_size=(6, 6)):
        # densities = torch.zeros((tile_size[-2] + 2) * (tile_size[-1] + 2) + 1)
        # tiles = ((self.channels.shape[-2] + tile_size[-2] - 1) // tile_size[-2]) * ((self.channels.shape[-1] + tile_size[-1] - 1) // tile_size[-1])
        densities = []

        if self.channels is None:
            return densities

        for y_off in range(0, self.channels.shape[-2], tile_size[-2]):
            for x_off in range(0, self.channels.shape[-1], tile_size[-1]):
                y_start = max(y_off - 1, 0)
                x_start = max(x_off - 1, 0)
                y_end = min(y_off + 1 + tile_size[0], self.channels.shape[-2])
                x_end = min(x_off + 1 + tile_size[1], self.channels.shape[-1])
                in_tile = self.channels[0, y_start:y_end, x_start:x_end]
                density = torch.sum(in_tile > 0)
                densities.append(density)

        return torch.tensor(densities)

    def get_mask(self):
        return self.channels

class ComputationsLogger(LoggingLayer):
    n_computed = 0
    n_samples = 0

    def __init__(self, **kwargs):
        super(ComputationsLogger, self).__init__(**kwargs)
        self.computed_mask = None
        self.n_samples = 0

    def reset(self):
        ComputationsLogger.n_computed = 0
        ComputationsLogger.n_samples = 0
        self.computed_mask = None
        self.n_samples = 0

    def forward(self, x):
        if not self.enabled:
            return x

        if not self.added_id:
            self.name = f"{LoggingLayer.id} {self.name}"
            LoggingLayer.id += 1
            self.added_id = True

        ComputationsLogger.n_computed += torch.sum(x)
        ComputationsLogger.n_samples += x.numel()

        self.n_samples += x.shape[0] * x.shape[1]
        comp_mask = torch.sum(torch.sum(x.clone().float(), dim=0), dim=0)
        if self.computed_mask is None or self.computed_mask.shape != x.shape[2:]:
            self.computed_mask = comp_mask
        else:
            self.computed_mask += comp_mask

        return x

    def get(self):
        if self.n_samples == 0:
            return 0
        return ComputationsLogger.n_computed / ComputationsLogger.n_samples

    def get_mask(self):
        if self.computed_mask is None:
            return None
        return self.computed_mask / self.n_samples

class MultiplicationsLogger(LoggingLayer):
    n_computed = 0
    n_samples = 0

    def __init__(self, **kwargs):
        super(MultiplicationsLogger, self).__init__(**kwargs)
        self.n_computed = 0
        self.n_samples = 0

    def reset(self):
        self.n_computed = 0
        self.n_samples = 0
        MultiplicationsLogger.n_computed = 0
        MultiplicationsLogger.n_samples = 0

    def forward(self, x, filter, in_mask=None, kernel_size=(3, 3), dilation=(1, 1), stride=(1, 1), padding=(0, 0), groups=1):
        if not self.added_id:
            self.name = f"{LoggingLayer.id} {self.name}"
            LoggingLayer.id += 1
            self.added_id = True

        if self.enabled:
            if in_mask is None:
                computed = torch.sum(x[:, 0]) * filter.numel()
                self.n_computed += computed
                MultiplicationsLogger.n_computed += computed
            else:
                n_filter = (filter[:, 0, 0, :] if filter.shape[1] == filter.shape[2] else filter[:, :, 0, 0]).numel()
                if filter.shape[2] == 1:
                    computed = torch.sum(x[:, 0]) * n_filter
                    self.n_computed += computed
                    MultiplicationsLogger.n_computed += computed
                else:
                    active_neighbors = count_active_neighbors(in_mask[:, :1], kernel_size=kernel_size, dilation=dilation, stride=stride, padding=padding)
                    active_neighbors = active_neighbors[:, 0][x[:, 0]]

                    # computed = torch.sum(active_neighbors) * filter[::groups, :, 0, 0].numel()
                    computed = torch.sum(active_neighbors) * n_filter
                    self.n_computed += computed
                    MultiplicationsLogger.n_computed += computed
            samples = x[:, 0].numel() * filter.numel()
            self.n_samples += samples
            MultiplicationsLogger.n_samples += samples

        return x

    def get_self_rel(self):
        if self.n_samples > 0:
            return self.n_computed / self.n_samples
        return 1

    def get_self_abs(self):
        return self.n_computed

    def get_global_rel(self):
        if MultiplicationsLogger.n_samples > 0:
            return MultiplicationsLogger.n_computed / MultiplicationsLogger.n_samples
        else:
            return 1

    def get_global_abs(self):
        return MultiplicationsLogger.n_computed

class PrevInBandwidthLogger(LoggingLayer):
    n_read = 0

    def __init__(self, **kwargs):
        super(PrevInBandwidthLogger, self).__init__(**kwargs)

    def reset(self):
        PrevInBandwidthLogger.n_read = 0

    def forward(self, mask: torch.Tensor, dilation: str, tile_size: int = 0):
        if not self.added_id:
            self.name = f"{LoggingLayer.id} {self.name}"
            LoggingLayer.id += 1
            self.added_id = True

        if self.enabled:
            if dilation == "tile":
                tile_size = tile_size if tile_size > 0 else mask.shape[2] // (-tile_size)
                mask = mask[:, :, ::tile_size, ::tile_size]
                inv_long_mask = (~mask).to(torch.long)
                inactive_neighbors = torch.zeros_like(inv_long_mask)
                inactive_neighbors[:, :, :-1] += inv_long_mask[:, :, 1:]
                inactive_neighbors[:, :, 1:] += inv_long_mask[:, :, :-1]
                inactive_neighbors[:, :, :, :-1] += inv_long_mask[:, :, :, 1:]
                inactive_neighbors[:, :, :, 1:] += inv_long_mask[:, :, :, :-1]
                inactive_neighbors[~mask] = 0
                n_read = torch.sum(inactive_neighbors).item() * tile_size
            else:
                inv_long_mask = (~mask).to(torch.long)
                inactive_neighbors = torch.zeros_like(inv_long_mask)
                inactive_neighbors[:, :, :-1] += inv_long_mask[:, :, 1:]
                inactive_neighbors[:, :, 1:] += inv_long_mask[:, :, :-1]
                inactive_neighbors[:, :, :, :-1] += inv_long_mask[:, :, :, 1:]
                inactive_neighbors[:, :, :, 1:] += inv_long_mask[:, :, :, :-1]
                inactive_neighbors[:, :, :-1, :-1] += inv_long_mask[:, :, 1:, 1:]
                inactive_neighbors[:, :, 1:, :-1] += inv_long_mask[:, :, :-1, 1:]
                inactive_neighbors[:, :, :-1, 1:] += inv_long_mask[:, :, 1:, :-1]
                inactive_neighbors[:, :, 1:, 1:] += inv_long_mask[:, :, :-1, :-1]
                inactive_neighbors[~mask] = 0
                n_read = torch.sum(inactive_neighbors).item()

            PrevInBandwidthLogger.n_read += n_read

        return mask

    def get(self):
        return PrevInBandwidthLogger.n_read

class PrevInputSizeLogger(LoggingLayer):
    n_vals = 0

    def __init__(self, **kwargs):
        super(PrevInputSizeLogger, self).__init__(**kwargs)
        self.n_vals = 0

    def reset(self):
        PrevInputSizeLogger.n_vals = 0
        self.n_vals = 0

    def forward(self, x):
        if not self.added_id:
            self.name = f"{LoggingLayer.id} {self.name}"
            LoggingLayer.id += 1
            self.added_id = True

        if x is None:
            return x

        PrevInputSizeLogger.n_vals += x[0].numel()
        self.n_vals += x[0].numel()

        return x

    def get_global(self):
        return PrevInputSizeLogger.n_vals

    def get(self):
        return self.n_vals

class PrevInputLogger(LoggingLayer):
    def __init__(self, **kwargs):
        super(PrevInputLogger, self).__init__(**kwargs)
        self.prev_input = None

    def reset(self):
        self.prev_input = None

    def forward(self, x):
        if not self.enabled:
            return x

        if not self.added_id:
            self.name = f"{LoggingLayer.id} {self.name}"
            LoggingLayer.id += 1
            self.added_id = True

        if x is None:
            return x

        prev_input = torch.sum(torch.sum(x.clone().float(), dim=0), dim=0)
        if self.prev_input is None or self.prev_input.shape != prev_input.shape:
            self.prev_input = prev_input
        else:
            self.prev_input += prev_input

        return x

    def get(self):
        if self.prev_input is None:
            return None
        return self.prev_input / torch.max(self.prev_input)

class InputLogger(PrevInputLogger):
    def __init__(self, **kwargs):
        super(InputLogger, self).__init__(**kwargs)

class OutputLogger(PrevInputLogger):
    def __init__(self, **kwargs):
        super(OutputLogger, self).__init__(**kwargs)

class RemoveDifferentialThreshold(LoggingLayer):
    def __init__(self, threshold=0.05, name="", enable_training=False, enabled=True, **kwargs):
        super(RemoveDifferentialThreshold, self).__init__(name=name, enabled=enabled)
        self.threshold = threshold
        self.prev_frame = None
        self.enable_training = enable_training

    def forward(self, x):
        if self.enabled and (torch.is_grad_enabled() and self.enable_training or not torch.is_grad_enabled()):
            if self.prev_frame is None:
                self.prev_frame = x.detach()
                return x
            mask = torch.abs(x - self.prev_frame) < self.threshold
            x[mask] = self.prev_frame[mask]
            self.prev_frame[~mask] = x[~mask].detach()

        return x

    def disable_logging(self):
        pass

    def enable_logging(self):
        pass

    def reset_history(self):
        self.prev_frame = None

class RemoveConsistentPixels(LoggingLayer):
    def __init__(self, threshold=0.05, name="", enable_training=False, enabled=True, **kwargs):
        super(RemoveConsistentPixels, self).__init__(name=name, enabled=enabled)
        self.threshold = threshold
        self.prev_frame = None
        self.enable_training = enable_training

    def forward(self, x):
        if self.enabled and (torch.is_grad_enabled() and self.enable_training or not torch.is_grad_enabled()):
            if self.prev_frame is None:
                self.prev_frame = x.detach()
                return x
            mask = torch.abs(x - self.prev_frame) < self.threshold
            x[mask] = 0.0
            self.prev_frame[~mask] = x[~mask].detach()

        return x

    def disable_logging(self):
        pass

    def enable_logging(self):
        pass

    def reset_history(self):
        self.prev_frame = None

class GradientLogFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, X, log_out):
        ctx.log_out = log_out
        return X

    @staticmethod
    def backward(ctx, grad):
        ctx.log_out['loss'] = grad
        return grad, None

class GradientLogger(LoggingLayer):
    def __init__(self, **kwargs):
        super(GradientLogger, self).__init__(**kwargs)
        self.fn = GradientLogFunction()
        self.loss = {'loss': None}

    def reset(self):
        self.loss['loss'] = None

    def forward(self, x):
        self.fn(x, self.loss)
        return x

    def get(self):
        return self.loss['loss']

class Loggers(LoggingLayer):
    threshold = 0.0

    def __init__(self, name="", modes=None, enabled=False, threshold=None, **kwargs):
        super(Loggers, self).__init__(name, enabled, **kwargs)
        self.threshold = threshold if threshold is not None else Loggers.threshold

        if modes is None:
            modes = ["DiffMask", "DiffDensity", "Density", "Mask"]
        self.logs = {}
        if "RemoveDiff" in modes:
            self.logs["RemoveDiff"] = RemoveDifferentialThreshold(enable_training=False, enabled=True, threshold=self.threshold)
        if "RemoveConsistent" in modes:
            self.logs["RemoveConsistent"] = RemoveConsistentPixels(enable_training=False, enabled=True, threshold=self.threshold)
        if "Density" in modes:
            self.logs["Density"] = DensityLogger(name=name, enabled=enabled, threshold=self.threshold, **kwargs)
        if "Mask" in modes:
            self.logs["Mask"] = PixelMaskLogger(name=name, enabled=enabled, threshold=self.threshold, **kwargs)
        if "DiffDensity" in modes:
            self.logs["DiffDensity"] = DifferentialDensityLogger(name=name, enabled=enabled, threshold=self.threshold, **kwargs)
        if "DiffMask" in modes:
            self.logs["DiffMask"] = DifferentialPixelMaskLogger(name=name, enabled=enabled, threshold=self.threshold, **kwargs)

        self.active_layers = ["ReLU", "Input"]
        self.always_use_removediff = False

    def reset(self):
        for name, log in self.logs.items():
            if issubclass(type(log), LoggingLayer):
                log.reset()

    def reset_history(self):
        for name, log in self.logs.items():
            if issubclass(type(log), LoggingLayer):
                log.reset_history()

    def _is_active(self):
        if self.active_layers is None:
            return True

        for layer in self.active_layers:
            if layer in self.name:
                return True

        return False

    def forward(self, x):
        if not self.added_id:
            self.name = f"{LoggingLayer.id} {self.name}"
            for name, log in list(self.logs.items()):
                self.logs[f"{LoggingLayer.id} {name}"] = log
                self.logs.pop(name)
            self.added_id = True

            for name, log in self.logs.items():
                log.name = name
                log.added_id = True

            LoggingLayer.id += 1

        if not self._is_active():
            if self.always_use_removediff:
                for name, log in self.logs.items():
                    if "RemoveDiff" in name:
                        x = log(x)
            return x

        for name, log in self.logs.items():
            if self.enabled:
                log.enable_logging()
            else:
                log.disable_logging()
            x = log(x)

        return x

    def get(self):
        res = {}

        if not self._is_active():
            return res

        for name, log in self.logs.items():
            if issubclass(type(log), LoggingLayer):
                res[name] = log.get()

        return res


import logging
import pathlib
import typing

from ._exceptions import ConfigurationPathsError

log = logging.getLogger(__name__)

PIPELINE_CONFIG_FILES = {
    "clients.yml",
    "release_states.yml",
    "deployments.yml",
    "frames.yml",
    "puff_map.yml",
    "service_principals.vault",
    "subscriptions.yml",
    "systems.yml",
    "tenants.yml",
}

T = typing.TypeVar("T", bound="PipelineConfigurationPaths")

class PipelineConfigurationPaths:
    """Paths to pipeline configuration files."""

    clients: pathlib.Path
    context: typing.Set[pathlib.Path]
    deployments: pathlib.Path
    frames: pathlib.Path
    puff_map: pathlib.Path
    release_states: pathlib.Path
    service_principals: pathlib.Path
    static_secrets: typing.Set[pathlib.Path]
    subscriptions: pathlib.Path
    systems: pathlib.Path
    tenants: pathlib.Path

    CONFIG_SUBDIRS: typing.Set[str] = {"static_secrets", "context"}

    def __init__(self: T) -> None:
        """Construct ``PipelineConfigurationPaths`` object."""
        for this_file in PIPELINE_CONFIG_FILES:
            path = pathlib.Path(this_file)
            setattr(self, path.stem, path)

    @classmethod
    def from_dict(cls: typing.Type[T], data: dict) -> T:
        """
        Construct ``PipelineConfigurationPaths`` object.

        NOTE: Delivering valid paths is the users responsibility.

        Args:
            data: Dictionary of data to populate in object.
        """
        this_object = cls()
        for x, y in data.items():
            setattr(this_object, x, y)

        return this_object

    @classmethod
    def from_paths(
        cls: typing.Type[T],
        client_config: pathlib.Path,
        system_config: pathlib.Path,
    ) -> T:
        """
        Construct ``PipelineConfigurationPaths`` object.

        Args:
            client_config: Path to client configuration directory.
            system_config: Path to system configuration directory.

        Raises:
            ConfigurationPathsError: If any paths are duplicated between
                client and system.
        """
        this_object = cls()
        client_files = this_object.__acquire_client_files(client_config)
        system_files = this_object.__acquire_system_files(system_config)

        if len(client_files + system_files) > len(PIPELINE_CONFIG_FILES):
            # must be duplicate files between the directories
            log.debug("client files, {0}".format(str(client_files)))
            log.debug("system files, {0}".format(str(system_files)))
            raise ConfigurationPathsError(
                "Duplicate files between "
                "directories, {0}, {1}".format(client_config, system_config)
            )

        this_object.static_secrets = cls.__acquire_static_secrets(client_config)
        this_object.context = cls.__acquire_template_context(
            client_config, system_config
        )

        return this_object

    @staticmethod
    def __acquire_static_secrets(
        client_config: pathlib.Path,
    ) -> typing.Set[pathlib.Path]:
        secrets_path = client_config / "static_secrets"
        result = PipelineConfigurationPaths.__acquire_subdir_files(
            secrets_path, "static secrets"
        )
        return result

    @staticmethod
    def __acquire_template_context(
        client_config: pathlib.Path, system_config: pathlib.Path
    ) -> typing.Set[pathlib.Path]:
        # template context could be located in either client or system config.
        context_client_path = client_config / "context"
        client_context_files = (
            PipelineConfigurationPaths.__acquire_subdir_files(
                context_client_path, "client template context"
            )
        )
        context_system_path = system_config / "context"
        system_context_files = (
            PipelineConfigurationPaths.__acquire_subdir_files(
                context_system_path, "system template context"
            )
        )
        result = client_context_files.union(system_context_files)
        return result

    def __acquire_client_files(
        self: T, client_config: pathlib.Path
    ) -> typing.List[pathlib.Path]:
        client_files = list()
        for x in client_config.iterdir():
            if (
                x.is_file()
                and (x.name in PIPELINE_CONFIG_FILES)
                and (x.stem not in self.CONFIG_SUBDIRS)
            ):
                log.info("adding client configuration file, {0}".format(x))
                setattr(self, x.stem, x)
                client_files.append(x)

        return client_files

    def __acquire_system_files(
        self: T, system_config: pathlib.Path
    ) -> typing.List[pathlib.Path]:
        system_files = list()
        for x in system_config.iterdir():
            if x.is_file() and (x.name in PIPELINE_CONFIG_FILES):
                log.info("adding system configuration file, {0}".format(x))
                setattr(self, x.stem, x)
                system_files.append(x)

        return system_files

    @staticmethod
    def __acquire_subdir_files(
        this_path: pathlib.Path, category: str
    ) -> typing.Set[pathlib.Path]:
        result = set()
        if this_path.is_dir():
            log.debug(
                "{2} directory, {0}, {1}".format(
                    this_path, str(list(this_path.iterdir())), category
                )
            )
            for x in this_path.iterdir():
                if x.is_file():
                    log.info(
                        "adding file to configuration, {1}, {0}".format(
                            x, category
                        )
                    )
                    result.add(x)
        elif this_path.exists():
            log.debug(f"{category} not a directory, {this_path}")
        else:
            log.debug(f"{category} not present, {this_path}")

        return result


from time import time
from random import randint

def get_size_table(*args) -> tuple:
    a1, b1, a2, b2 = args
    squares = []
    laptop_sizes = [(max(a1, b1), min(a1, b1)), (max(a2, b2), min(a2, b2))]
    laptop_sizes.sort(key=lambda x: x[0], reverse=True)

    length_1, width_1 = laptop_sizes[0][0] + laptop_sizes[1][0], max((laptop_sizes[0][1], laptop_sizes[1][1]))
    length_2, width_2 = laptop_sizes[0][0], laptop_sizes[0][1] + laptop_sizes[1][1]
    length_3, width_3 = laptop_sizes[0][0] + laptop_sizes[1][1], max((laptop_sizes[0][1], laptop_sizes[1][0]))

    squares.append(length_1 * width_1)
    squares.append(length_2 * width_2)
    squares.append(length_3 * width_3)
    min_square = min(squares)

    if squares[0] == min_square:
        return length_1, width_1
    elif squares[1] == min_square:
        return length_2, width_2
    return length_3, width_3

def main():
    a1, b1, a2, b2 = (randint(1, 1000) for _ in range(4))
    print(a1, b1, a2, b2)
    print(*get_size_table(a1, b1, a2, b2))

def test():
    assert get_size_table(10, 2, 2, 10) == (20, 2) or (2, 20) or (4, 10) or (10, 4)
    assert get_size_table(5, 7, 3, 2) == (9, 5) or (5, 9)

if __name__ == '__main__':
    time_st = time()
    main()
    print(time() - time_st)