import random
import typing as t

import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from langchain.embeddings import HuggingFaceInstructEmbeddings
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.preprocessing import MinMaxScaler

from bunkatopics.datamodel import BourdieuDimension, ContinuumDimension, Document, Term
from bunkatopics.functions.topic_document import get_top_documents
from bunkatopics.functions.topic_gen_representation import get_clean_topic_all
from bunkatopics.functions.topics_modeling import get_topics
from bunkatopics.visualisation.explainer import plot_specific_terms
from bunkatopics.visualisation.visu_utils import wrap_by_word

pd.options.mode.chained_assignment = None

def get_continuum(
    embedding_model: HuggingFaceInstructEmbeddings,
    docs: t.List[Document],
    cont_name: str = "emotion",
    left_words: list = ["hate", "pain"],
    right_words: list = ["love", "good"],
    scale: bool = False,
) -> t.List[Document]:
    df_docs = pd.DataFrame.from_records([doc.dict() for doc in docs])
    df_emb = df_docs[["doc_id", "embedding"]]
    df_emb = df_emb.set_index("doc_id")
    df_emb = pd.DataFrame(list(df_emb["embedding"]))
    df_emb.index = df_docs["doc_id"]

    continuum = ContinuumDimension(
        id=cont_name, left_words=left_words, right_words=right_words
    )

    # Compute the extremity embeddings
    left_embedding = embedding_model.embed_documents(continuum.left_words)
    right_embedding = embedding_model.embed_documents(continuum.right_words)

    left_embedding = pd.DataFrame(left_embedding).mean().values.reshape(1, -1)
    right_embedding = pd.DataFrame(right_embedding).mean().values.reshape(1, -1)

    # Make the difference to get the continnum
    continuum_embedding = left_embedding - right_embedding
    df_continuum = pd.DataFrame(continuum_embedding)
    df_continuum.index = ["distance"]

    # Compute the Cosine Similarity
    full_emb = pd.concat([df_emb, df_continuum])
    df_bert = pd.DataFrame(cosine_similarity(full_emb))

    df_bert.index = full_emb.index
    df_bert.columns = full_emb.index
    df_bert = df_bert.iloc[-1:,].T
    df_bert = df_bert.sort_values("distance", ascending=False).reset_index()
    df_bert = df_bert[1:]
    df_bert = df_bert.rename(columns={"index": "doc_id"})
    final_df = pd.merge(df_bert, df_docs[["doc_id", "content"]], on="doc_id")

    if scale:
        scaler = MinMaxScaler(feature_range=(-1, 1))
        final_df[["distance"]] = scaler.fit_transform(final_df[["distance"]])

    final_df = final_df.set_index("doc_id")
    final_df = final_df[["distance"]]

    distance_dict = final_df.to_dict("index")

    new_docs = docs.copy()

    for doc in new_docs:
        res = BourdieuDimension(
            continuum=continuum, distance=distance_dict.get(doc.doc_id)["distance"]
        )
        doc.bourdieu_dimensions.append(res)

    return new_docs

def plot_unique_dimension(
    docs: t.List[Document],
    id: str = id,
    left: list = ["aggressivity"],
    right: list = ["peacefullness"],
    height=700,
    width=600,
    explainer: bool = True,
    explainer_ngrams: list = [1, 2],
) -> go.Figure:
    left = " ".join(left)
    right = " ".join(right)

    distances = [
        x.distance
        for doc in docs
        for x in doc.bourdieu_dimensions
        if x.continuum.id == id
    ]
    doc_id = [x.doc_id for x in docs]
    content = [x.content for x in docs]

    df_distances = pd.DataFrame(
        {"doc_id": doc_id, "distances": distances, "content": content}
    )

    name = "<" + right + "-" + left + ">"

    df_fig = df_distances.rename(columns={"distances": name})
    df_fig["content"] = df_fig["content"].apply(lambda x: wrap_by_word(x, 10))

    fig = px.box(
        df_fig,
        y=name,
        points="all",
        hover_data=["content"],
        height=height,
        width=width,
        template="plotly_white",
    )

    fig.add_shape(
        dict(
            type="line",
            x0=df_fig[name].min(),  # Set the minimum x-coordinate of the line
            x1=df_fig[name].max(),  # Set the maximum x-coordinate of the line
            y0=0,
            y1=0,
            line=dict(color="red", width=4),
        )
    )
    if explainer:
        plot_specific_terms(
            docs=docs,
            left_words=left,
            right_words=right,
            id=id,
            ngrams=explainer_ngrams,
            quantile=0.80,
            top_n=20,
        )
    return fig

def visualize_bourdieu_one_dimension(
    docs: t.List[Document],
    embedding_model,
    left: str = ["aggressivity"],
    right: str = ["peacefullness"],
    height=700,
    width=600,
    explainer: bool = True,
    explainer_ngrams: list = [1, 2],
) -> go.Figure:
    id = str(random.randint(0, 10000))

    new_docs = get_continuum(
        embedding_model=embedding_model,
        docs=docs,
        cont_name=id,
        left_words=left,
        right_words=right,
        scale=False,
    )

    fig = plot_unique_dimension(
        new_docs,
        id=id,
        left=left,
        right=right,
        height=height,
        width=width,
        explainer=explainer,
        explainer_ngrams=explainer_ngrams,
    )
    return fig

def visualize_bourdieu(
    embedding_model,
    generative_model,
    docs: t.List[Document],
    terms: t.List[Term],
    x_left_words: t.List[str] = ["war"],
    x_right_words: t.List[str] = ["peace"],
    y_top_words: t.List[str] = ["men"],
    y_bottom_words: t.List[str] = ["women"],
    height: int = 1500,
    width: int = 1500,
    clustering: bool = True,
    topic_gen_name: bool = False,
    topic_n_clusters: int = 5,
    topic_terms: int = 2,
    topic_ngrams: list = [1, 2],
    display_percent: bool = True,
    use_doc_gen_topic: bool = False,
    gen_topic_language: str = "english",
    label_size_ratio_label: int = 50,
    topic_top_terms_overall: int = 500,
    manual_axis_name: dict = None,
    radius_size: float = 0.3,
    convex_hull: bool = True,
):
    # Reset
    for doc in docs:
        doc.bourdieu_dimensions = []

    # Compute Continuums
    new_docs = get_continuum(
        embedding_model,
        docs,
        cont_name="cont1",
        left_words=x_left_words,
        right_words=x_right_words,
    )
    new_docs = get_continuum(
        embedding_model,
        docs,
        cont_name="cont2",
        left_words=y_top_words,
        right_words=y_bottom_words,
    )

    df_names = [
        {
            "names": [y.continuum.id for y in x.bourdieu_dimensions],
            "left_words": [y.continuum.left_words for y in x.bourdieu_dimensions],
            "right_words": [y.continuum.right_words for y in x.bourdieu_dimensions],
        }
        for x in new_docs
    ]
    df_names = pd.DataFrame(df_names)
    df_names = df_names.explode(["names", "left_words", "right_words"])
    df_names["left_words"] = df_names["left_words"].apply(lambda x: "-".join(x))
    df_names["right_words"] = df_names["right_words"].apply(lambda x: "-".join(x))
    df_names = df_names.drop_duplicates()
    df_names = df_names.set_index("names")

    dict_bourdieu = df_names.to_dict(orient="index")

    df_bourdieu = [
        {
            "doc_id": x.doc_id,
            "coordinates": [y.distance for y in x.bourdieu_dimensions],
            "names": [y.continuum.id for y in x.bourdieu_dimensions],
        }
        for x in new_docs
    ]

    df_bourdieu = pd.DataFrame(df_bourdieu)
    df_bourdieu = df_bourdieu.explode(["coordinates", "names"])

    # Filter with only the top and bottom data to avoid getting results too far form the continnuums

    df_content = [{"doc_id": x.doc_id, "content": x.content} for x in new_docs]
    df_content = pd.DataFrame(df_content)

    df_fig = df_bourdieu[["doc_id", "coordinates", "names"]]
    df_fig = df_fig.pivot(index="doc_id", columns="names", values="coordinates")
    df_fig = df_fig.reset_index()

    # Remove the data inside the radius of 1/3 of max because central data does not mean mucj
    df_fig["cont1"] = df_fig["cont1"].astype(float)
    df_fig["cont2"] = df_fig["cont2"].astype(float)

    import numpy as np

    x_values = df_fig["cont1"].values
    y_values = df_fig["cont2"].values

    distances = np.sqrt(x_values**2 + y_values**2)

    circle_radius = max(df_fig.cont1) * radius_size

    df_fig["distances"] = distances
    df_fig["outside"] = "0"
    df_fig["outside"][df_fig["distances"] >= circle_radius] = "1"

    outside_ids = list(df_fig["doc_id"][df_fig["outside"] == "1"])

    df_fig = df_fig[df_fig["doc_id"].isin(outside_ids)]

    df_fig = pd.merge(df_content, df_fig, on="doc_id")
    df_fig["Text"] = df_fig["content"].apply(lambda x: wrap_by_word(x, 10))

    x_axis_name = list(dict_bourdieu.keys())[0]
    y_axis_name = list(dict_bourdieu.keys())[1]

    x_left_words = dict_bourdieu[x_axis_name]["left_words"]
    x_right_words = dict_bourdieu[x_axis_name]["right_words"]
    y_top_words = dict_bourdieu[y_axis_name]["left_words"]
    y_bottom_words = dict_bourdieu[y_axis_name]["right_words"]

    fig = go.Figure(
        go.Histogram2dContour(
            x=df_fig[x_axis_name],
            y=df_fig[y_axis_name],
            colorscale="delta",
            showscale=False,
        ),
    )

    scatter_fig = px.scatter(
        df_fig,
        x=x_axis_name,
        y=y_axis_name,
        color="outside",
        color_discrete_map={"1": "white", "0": "grey"},
        hover_data=["Text"],
        template="simple_white",
        height=height,
        width=width,
        opacity=0.3,
        # title="Bourdieu Plot"
        # color_discrete_sequence=["blue"],
    )

    for trace in scatter_fig.data:
        fig.add_trace(trace)

    # Set the axis to the max value to get a square
    max_val = max(
        abs(min(df_fig[y_axis_name])),
        abs(max(df_fig[y_axis_name])),
        abs(max(df_fig[x_axis_name])),
        abs(min(df_fig[x_axis_name])),
    )

    # Add axis lines for x=0 and y=0
    fig.add_shape(
        type="line",
        x0=0,
        x1=0,
        # y0=-max_val,
        # y1=max_val,
        y0=min(df_fig[y_axis_name]),
        y1=max(df_fig[y_axis_name]),
        line=dict(color="white", width=3),  # Customize line color and width
    )

    fig.add_shape(
        type="line",
        x0=min(df_fig[x_axis_name]),
        x1=max(df_fig[x_axis_name]),
        # x0=-max_val,
        # x1=max_val,
        y0=0,
        y1=0,
        line=dict(color="white", width=3),  # Customize line color and width
    )

    fig.update_layout(
        font_size=25,
        width=width,
        height=height,
        margin=dict(
            t=width / 50,
            b=width / 50,
            r=width / 50,
            l=width / 50,
        ),
        # title=dict(font=dict(size=width / 40)),
    )

    fig.update_layout(showlegend=False)
    """
    histogram2d_contour = go.Figure(
        go.Histogram2dContour(
            x=df_fig[x_axis_name],
            y=df_fig[y_axis_name],
            colorscale="delta",
            showscale=False,
        ),
    )
    fig.add_trace(histogram2d_contour.data[0])

    scatter_fig = px.scatter(
        df_fig,
        x=x_axis_name,
        y=y_axis_name,
        color="outside",
        color_discrete_map={"1": "white", "0": "grey"},
        hover_data=["Text"],
        template="simple_white",
        height=height,
        width=width,
        opacity=0.3,
        # title="Bourdieu Plot"
        # color_discrete_sequence=["blue"],
    )

    for trace in scatter_fig.data:
        fig.add_trace(trace)

    """

    """

    fig.update_xaxes(
        showgrid=False,
        showticklabels=False,
        zeroline=True,
        zerolinecolor="white",
        zerolinewidth=2,
    )
    fig.update_yaxes(
        showgrid=False,
        showticklabels=False,
        zeroline=True,
        zerolinecolor="white",
        zerolinewidth=2,
    )

    """

    if manual_axis_name is None:
        y_top_name = y_top_words
        y_bottom_name = y_bottom_words
        x_left_name = x_left_words
        x_right_name = x_right_words

    else:
        y_top_name = manual_axis_name["y_top_name"]
        y_bottom_name = manual_axis_name["y_bottom_name"]
        x_left_name = manual_axis_name["x_left_name"]
        x_right_name = manual_axis_name["x_right_name"]

    fig.update_layout(
        annotations=[
            dict(
                x=0,
                # y=max_val,
                y=max(df_fig[y_axis_name]),
                xref="x",
                yref="y",
                text=y_top_name,
                showarrow=False,
                xanchor="right",
                yanchor="top",
                font=dict(size=width / label_size_ratio_label, color="white"),
            ),
            dict(
                x=0,
                y=min(df_fig[y_axis_name]),
                # y=-max_val,
                xref="x",
                yref="y",
                text=y_bottom_name,
                showarrow=False,
                xanchor="left",
                yanchor="bottom",
                font=dict(size=width / label_size_ratio_label, color="white"),
            ),
            dict(
                x=max(df_fig[x_axis_name]),
                # x=max_val,
                y=0,
                xref="x",
                yref="y",
                text=x_left_name,
                showarrow=False,
                xanchor="right",
                yanchor="top",
                font=dict(size=width / label_size_ratio_label, color="white"),
            ),
            dict(
                x=min(df_fig[x_axis_name]),
                # x=-max_val,
                y=0,
                xref="x",
                yref="y",
                text=x_right_name,
                showarrow=False,
                xanchor="left",
                yanchor="bottom",
                font=dict(size=width / label_size_ratio_label, color="white"),
            ),
        ]
    )

    if clustering:
        df_bourdieu_pivot = df_bourdieu.pivot(
            index="doc_id", columns="names", values="coordinates"
        )
        df_bourdieu_pivot = df_bourdieu_pivot.reset_index()
        df_bourdieu_pivot.columns = ["doc_id", "x", "y"]
        df_bourdieu_pivot = df_bourdieu_pivot.set_index("doc_id")

        dict_doc = df_bourdieu_pivot[["x", "y"]].to_dict("index")

        for doc in new_docs:
            doc.x = dict_doc.get(doc.doc_id)["x"]
            doc.y = dict_doc.get(doc.doc_id)["y"]

        new_docs = [doc for doc in new_docs if doc.doc_id in outside_ids]

        bourdieu_topics = get_topics(
            docs=new_docs,
            terms=terms,
            n_clusters=topic_n_clusters,
            ngrams=topic_ngrams,
            name_lenght=topic_terms,
            top_terms_overall=topic_top_terms_overall,
        )

        if topic_gen_name:
            # Get top documents for the generative AI query
            new_docs = get_top_documents(new_docs, bourdieu_topics, ranking_terms=20)

            bourdieu_topics = get_clean_topic_all(
                generative_model,
                language=gen_topic_language,
                topics=bourdieu_topics,
                docs=new_docs,
                use_doc=use_doc_gen_topic,
            )

        label_size_ratio_clusters = 100
        topics_x = [x.x_centroid for x in bourdieu_topics]
        topics_y = [x.y_centroid for x in bourdieu_topics]
        topic_names = [x.name for x in bourdieu_topics]
        topics_name_plotly = [wrap_by_word(x, 7) for x in topic_names]

        # Display Topics
        for x, y, label in zip(topics_x, topics_y, topics_name_plotly):
            fig.add_annotation(
                x=x,
                y=y,
                text=label,
                font=dict(
                    family="Courier New, monospace",
                    size=width / label_size_ratio_clusters,
                    color="red",
                ),
                bordercolor="#c7c7c7",
                borderwidth=width / 1000,
                borderpad=width / 500,
                bgcolor="white",
                opacity=1,
            )

        if convex_hull:
            try:
                for topic in bourdieu_topics:
                    # Create a Scatter plot with the convex hull coordinates
                    trace = go.Scatter(
                        x=topic.convex_hull.x_coordinates,
                        y=topic.convex_hull.y_coordinates,  # Assuming y=0 for simplicity
                        mode="lines",
                        name="Convex Hull",
                        line=dict(color="grey"),
                        showlegend=False,
                    )

                    fig.add_trace(trace)
            except:
                pass

    if display_percent:
        # Calculate the percentage for every box

        df_fig_percent = df_fig[df_fig["doc_id"].isin(outside_ids)]

        label_size_ratio_percent = 20
        opacity = 0.4
        case1_count = len(
            df_fig_percent[
                (df_fig_percent["cont1"] < 0) & (df_fig_percent["cont2"] < 0)
            ]
        )
        total_count = len(df_fig_percent)
        case1_percentage = str(round((case1_count / total_count) * 100, 1)) + "%"

        fig.add_annotation(
            x=min(df_fig_percent[x_axis_name]),
            y=min(df_fig_percent[y_axis_name]),
            text=case1_percentage,
            font=dict(
                family="Courier New, monospace",
                size=width / label_size_ratio_percent,
                color="grey",
            ),
            opacity=opacity,
            xanchor="left",
        )

        case2_count = len(
            df_fig_percent[
                (df_fig_percent["cont1"] < 0) & (df_fig_percent["cont2"] > 0)
            ]
        )
        case2_percentage = str(round((case2_count / total_count) * 100, 1)) + "%"

        fig.add_annotation(
            x=min(df_fig_percent[x_axis_name]),
            y=max(df_fig_percent[y_axis_name]),
            text=case2_percentage,
            font=dict(
                family="Courier New, monospace",
                size=width / label_size_ratio_percent,
                color="grey",
            ),
            opacity=opacity,
            xanchor="left",
        )

        case3_count = len(
            df_fig_percent[
                (df_fig_percent["cont1"] > 0) & (df_fig_percent["cont2"] < 0)
            ]
        )
        case3_percentage = str(round((case3_count / total_count) * 100, 1)) + "%"

        fig.add_annotation(
            x=max(df_fig_percent[x_axis_name]),
            y=min(df_fig_percent[y_axis_name]),
            text=case3_percentage,
            font=dict(
                family="Courier New, monospace",
                size=width / label_size_ratio_percent,
                color="grey",
            ),
            opacity=opacity,
            xanchor="left",
        )

        case4_count = len(
            df_fig_percent[
                (df_fig_percent["cont1"] > 0) & (df_fig_percent["cont2"] > 0)
            ]
        )
        case4_percentage = str(round((case4_count / total_count) * 100, 1)) + "%"

        fig.add_annotation(
            x=max(df_fig_percent[x_axis_name]),
            y=max(df_fig_percent[y_axis_name]),
            text=case4_percentage,
            font=dict(
                family="Courier New, monospace",
                size=width / label_size_ratio_percent,
                color="grey",
            ),
            opacity=opacity,
            xanchor="left",
        )

    # Update the x-axis and y-axis labels
    fig.update_xaxes(
        title_text="",
        scaleanchor="y",
        scaleratio=1,
        showgrid=False,
        showticklabels=False,
        zeroline=True,
        zerolinecolor="white",
        zerolinewidth=2,
    )
    fig.update_yaxes(
        title_text="",
        scaleanchor="x",
        scaleratio=1,
        showgrid=False,
        showticklabels=False,
        zeroline=True,
        zerolinecolor="white",
        zerolinewidth=2,
    )

    return fig, df_bourdieu


# coding=utf-8
#
# /**************************************************************************
# ***
# ***	File Author: Dell, 2018年 09月 18日 星期二 16:28:12 CST
# ***
# **************************************************************************/
#

import os
import sys
import logging
import argparse

import model

parser = argparse.ArgumentParser(description='Train Image Classificer Model')

parser.add_argument(
    '-root-dir',
    type=str,
    default=model.DEFAULT_TRAIN_DATA_ROOT_DIR,
    help='train data root directory, default: ' +
    model.DEFAULT_TRAIN_DATA_ROOT_DIR)
parser.add_argument(
    '-epochs',
    type=int,
    default=32,
    help='number of epochs for train, default: 32')
parser.add_argument(
    '-batch-size',
    type=int,
    default=64,
    help='batch size for training, default: 64')
parser.add_argument(
    '-device',
    type=str,
    default="cuda:0",
    help='cuda:0 or cpu, default: cuda:0')

def makedirs():
    for d in ["logs", "model"]:
        if not os.path.exists(d):
            os.mkdir(d)
        if not os.path.isdir(d):
            logging.error(
                "Please create dir 'logs' or 'model' under current directory.")
            raise Exception("logs or model is not directory.")

if __name__ == '__main__':
    args = parser.parse_args()

    if (not os.path.exists(args.root_dir)) or (not os.path.isdir(
            args.root_dir)):
        logging.error(args.root_dir + ' is not director or not exists.')
        sys.exit(-1)

    makedirs()

    data = model.train_data_loader(args.root_dir, args.batch_size)
    net = model.load_model(args.device, model.DEFAULT_MODEL)
    model.train_model(args.device, net, data, args.epochs)


# by filtering stock that is in the range of 0.5 to 2 pct difference

import os
from dotenv import load_dotenv
load_dotenv()
import os
from supabase import create_client
import numpy as np
import pandas as pd
import requests
from datetime import datetime
from io import StringIO

def preprocess_numeric_value(value):
    if pd.isna(value) or value == 0.0:
        return np.nan
    str_value = str(value)
    if 'T' in str_value.upper():
        return float(str_value.upper().replace('T', '')) * 1e12
    elif 'B' in str_value.upper():
        return float(str_value.upper().replace('B', '')) * 1e9
    elif 'M' in str_value.upper():
        return float(str_value.upper().replace('M', '')) * 1e6
    elif 'K' in str_value.upper():
        return float(str_value.upper().replace('K', '')) * 1e3
    else:
        return float(value)

def preprocess_percentage_value(value):
    if pd.isna(value):
        return np.nan
    if '%' in str(value):
        return float(str(value).replace('%', '').replace(',', ''))/100
    else:
        return float(str(value))

def calculate_growth(row, y, type=None):
    try:
        year2 = row[y]

        if type == 'revenue':
            year1 = row['total_revenue']/row['multiplier']
        else:
            year1 = row['basic_eps']

        if pd.isna(year1) or year1 == 0.0:
            return np.nan
        else:
            return (year2 - year1) / year1
    except (ValueError, KeyError):
        return np.nan

url = os.environ.get("SUPABASE_URL")
key = os.environ.get("SUPABASE_KEY")
supabase = create_client(url, key)

key_data = supabase.table("idx_company_profile").select("symbol","sub_sector_id").execute()
key_df = pd.DataFrame(key_data.data).sort_values(['symbol'])

symbols = key_df['symbol'].to_list()

headers = {'User-Agent': 'Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:109.0) Gecko/20100101 Firefox/113.0'}

all_list = {
    'revenue_year_ago' : [],
    'avg_estimate_earnings_current_year': [],
    'avg_estimate_earnings_next_year': [],
    'avg_estimate_revenue_current_year': [],
    'avg_estimate_revenue_next_year': [],
    'estimate_overall_growth_current_year': [],
    'estimate_overall_growth_next_year': [],
    'estimate_overall_growth_next_five_years': [],
}

for symbol in symbols:
    try:
        url = f'https://finance.yahoo.com/quote/{symbol}/analysis?p={symbol}'
        html_content = requests.get(url, headers=headers).text
        source_df = pd.read_html(StringIO(html_content))

        earnings_df = source_df[0]
        revenue_df = source_df[1]
        overall_growth_df = source_df[5]

        avg_estimate_earnings_current_year = earnings_df.loc[earnings_df['Earnings Estimate'] == 'Avg. Estimate'].iloc[:, 3].values[0]
        avg_estimate_earnings_next_year = earnings_df.loc[earnings_df['Earnings Estimate'] == 'Avg. Estimate'].iloc[:, 4].values[0]
        avg_estimate_revenue_current_year = revenue_df.loc[revenue_df['Revenue Estimate'] == 'Avg. Estimate'].iloc[:, 3].values[0]
        year_ago = revenue_df.loc[revenue_df['Revenue Estimate'] == 'Year Ago Sales'].iloc[:, 3].values[0]
        avg_estimate_revenue_next_year = revenue_df.loc[revenue_df['Revenue Estimate'] == 'Avg. Estimate'].iloc[:, 4].values[0]
        estimate_overall_growth_current_year = overall_growth_df.loc[overall_growth_df['Growth Estimates'] == 'Current Year'].iloc[0, :].values[1]
        estimate_overall_growth_next_year = overall_growth_df.loc[overall_growth_df['Growth Estimates'] == 'Next Year'].iloc[0, :].values[1]
        estimate_overall_growth_next_five_years = overall_growth_df.loc[overall_growth_df['Growth Estimates'] == 'Next 5 Years (per annum)'].iloc[0, :].values[1]

        all_list['avg_estimate_earnings_current_year'].append(avg_estimate_earnings_current_year)
        all_list['avg_estimate_earnings_next_year'].append(avg_estimate_earnings_next_year)
        all_list['avg_estimate_revenue_current_year'].append(avg_estimate_revenue_current_year)
        all_list['avg_estimate_revenue_next_year'].append(avg_estimate_revenue_next_year)
        all_list['revenue_year_ago'].append(year_ago)
        all_list['estimate_overall_growth_current_year'].append(estimate_overall_growth_current_year)
        all_list['estimate_overall_growth_next_year'].append(estimate_overall_growth_next_year)
        all_list['estimate_overall_growth_next_five_years'].append(estimate_overall_growth_next_five_years)

        print(f"{symbol} data processed")
    except Exception as e:
        for key in all_list.keys():
            all_list[key].append(np.nan)
        print(f"{symbol} no data")

data_dict = {
    'symbol': symbols,
    **all_list,
}

forecast_df = pd.DataFrame.from_dict(data_dict)

current_year = datetime.now().year
last_year= f"{current_year-1}-12-31"
db_data = supabase.table("idx_financials_annual").select("symbol","total_revenue","basic_eps").eq("date", last_year).execute()
db_df = pd.DataFrame(db_data.data).sort_values(['symbol'])

df = forecast_df.merge(db_df, on='symbol', how='inner').merge(key_df, on='symbol', how='inner')

numeric_columns = ['avg_estimate_earnings_current_year', 'avg_estimate_earnings_next_year', 'avg_estimate_revenue_current_year', 'avg_estimate_revenue_next_year','revenue_year_ago']

for column in numeric_columns:
    df[column] = df[column].apply(preprocess_numeric_value)

percentage_columns = ['estimate_overall_growth_current_year', 'estimate_overall_growth_next_year',	'estimate_overall_growth_next_five_years']
for percentage_column in percentage_columns:
  df[str(percentage_column)] = df[str(percentage_column)].apply(preprocess_percentage_value)

df['multiplier'] = 1
df_1000 = df.copy()
df_1000['multiplier'] = 1000
df_1000['revenue_year_ago'] = df_1000['revenue_year_ago'] * df_1000['multiplier']
growth_forecast_df = pd.concat([df, df_1000], axis=0, ignore_index=True)
growth_forecast_df = growth_forecast_df.sort_values(by=["symbol", "multiplier"])
growth_forecast_df['ratio_mult'] = growth_forecast_df['total_revenue']/ growth_forecast_df['revenue_year_ago']
growth_forecast_df = growth_forecast_df.query("ratio_mult > 0.5 and ratio_mult < 2")
growth_forecast_df.to_csv('idx_company_rev_year_ago_filtered.csv',index = False)

numeric_columns = ['avg_estimate_earnings_current_year', 'avg_estimate_earnings_next_year', 'avg_estimate_revenue_current_year', 'avg_estimate_revenue_next_year']

for column in numeric_columns:
    type = None
    if 'revenue' in column:
        type = 'revenue'
    growth_forecast_df[f'{column[4:12]}_growth_{column[13:]}'] = growth_forecast_df.apply(calculate_growth, y=column, type=type, axis=1)

final_df = growth_forecast_df[['symbol','sub_sector_id','estimate_overall_growth_current_year','estimate_overall_growth_next_year','estimate_overall_growth_next_five_years','avg_estimate_earnings_current_year','avg_estimate_earnings_next_year','estimate_growth_earnings_current_year','estimate_growth_earnings_next_year','avg_estimate_revenue_current_year','avg_estimate_revenue_next_year','estimate_growth_revenue_current_year','estimate_growth_revenue_next_year']]
final_df.columns = ['symbol','sub_sector_id','overall_growth_current_year_f','overall_growth_next_year_f','overall_growth_next_five_years_f','avg_eps_current_year','avg_eps_next_year','eps_growth_current_year_f','eps_growth_next_year_f','avg_revenue_current_year','avg_revenue_next_year','revenue_growth_current_year_f','revenue_growth_next_year_f']
final_df.to_csv('idx_company_growth_forecast.csv',index = False)
try:
    result = supabase.table("idx_company_growth_forecast").upsert(final_df.to_dict(orient='records'), returning='minimal', on_conflict=['symbol'])
    print("Upsert operation successful.")
except Exception as e:
    print(f"Error during upsert operation: {e}")


"""
Serialize data to/from JSON
"""

# Avoid shadowing the standard library json module
from __future__ import absolute_import
from __future__ import unicode_literals

import datetime
import decimal
import json
import sys

from theory.core.serializers.base import DeserializationError
from theory.core.serializers.python import Serializer as PythonSerializer
from theory.core.serializers.python import Deserializer as PythonDeserializer
from theory.utils import six
from theory.utils.timezone import isAware

class Serializer(PythonSerializer):
  """
  Convert a queryset to JSON.
  """
  internalUseOnly = False

  def startSerialization(self):
    if json.__version__.split('.') >= ['2', '1', '3']:
      # Use JS strings to represent Python Decimal instances (ticket #16850)
      self.options.update({'useDecimal': False})
    self._current = None
    self.jsonKwargs = self.options.copy()
    self.jsonKwargs.pop('stream', None)
    self.jsonKwargs.pop('fields', None)
    if self.options.get('indent'):
      # Prevent trailing spaces
      self.jsonKwargs['separators'] = (',', ': ')
    self.stream.write("[")

  def endSerialization(self):
    if self.options.get("indent"):
      self.stream.write("\n")
    self.stream.write("]")
    if self.options.get("indent"):
      self.stream.write("\n")

  def endObject(self, obj):
    # self._current has the field data
    indent = self.options.get("indent")
    if not self.first:
      self.stream.write(",")
      if not indent:
        self.stream.write(" ")
    if indent:
      self.stream.write("\n")
    json.dump(self.getDumpObject(obj), self.stream,
         cls=TheoryJSONEncoder, **self.jsonKwargs)
    self._current = None

  def getvalue(self):
    # Grand-parent super
    return super(PythonSerializer, self).getvalue()

def Deserializer(streamOrString, **options):
  """
  Deserialize a stream or string of JSON data.
  """
  if not isinstance(streamOrString, (bytes, six.stringTypes)):
    streamOrString = streamOrString.read()
  if isinstance(streamOrString, bytes):
    streamOrString = streamOrString.decode('utf-8')
  try:
    objects = json.loads(streamOrString)
    for obj in PythonDeserializer(objects, **options):
      yield obj
  except GeneratorExit:
    raise
  except Exception as e:
    # Map to deserializer error
    six.reraise(DeserializationError, DeserializationError(e), sys.exc_info()[2])

class TheoryJSONEncoder(json.JSONEncoder):
  """
  JSONEncoder subclass that knows how to encode date/time and decimal types.
  """
  def default(self, o):
    # See "Date Time String Format" in the ECMA-262 specification.
    if isinstance(o, datetime.datetime):
      r = o.isoformat()
      if o.microsecond:
        r = r[:23] + r[26:]
      if r.endswith('+00:00'):
        r = r[:-6] + 'Z'
      return r
    elif isinstance(o, datetime.date):
      return o.isoformat()
    elif isinstance(o, datetime.time):
      if isAware(o):
        raise ValueError("JSON can't represent timezone-aware times.")
      r = o.isoformat()
      if o.microsecond:
        r = r[:12]
      return r
    elif isinstance(o, decimal.Decimal):
      return str(o)
    else:
      return super(TheoryJSONEncoder, self).default(o)

# Older, deprecated class name (for backwards compatibility purposes).
DateTimeAwareJSONEncoder = TheoryJSONEncoder


from flask import Flask
from flask import request, jsonify
import json, os, util, pickle
app = Flask(__name__)
SIMULATION_RESULT_PATH = './sim_result'

from flask_cors import CORS
CORS(app)

def load(path):
    with open(path, 'rb') as f:
        obj = pickle.load(f)
    return obj

@app.route("/")
def hello_world():
    return "
InterSim Beta Server is running.
"

@app.route('/vis', methods=['GET', 'POST'])
def visualization_get():
    sim_name = request.args.get('sim')
    scene_id = request.args.get('sceneid')
    file_id = request.args.get('fileid')
    if request.method == 'GET':
        with open(f"sim_result/{sim_name}/json/{file_id}.json", "r") as json_file:
            my_dict = json.load(json_file)
            keys = list(my_dict.keys())
            my_dict['selected_scene_index'] = keys.index(scene_id)
        return jsonify(my_dict)

@app.route('/list_scenarios')
def get_simulations():
    return summary_simulations()

@app.route('/list_simulation_selection')
def get_simulation_selection():
    dataset = request.args.get('dataset')
    return list_simulation_selection(dataset=dataset)

@app.route('/list_senarios')
def get_scenarios_list():
    simulation_name = request.args.get('simulation')
    return list_scenarios(simulation_name)

def check_path_valid(path):
    if not os.path.isdir(path):
        return False
    if not os.path.exists(os.path.join(path, 'sim.info')):
        return False
    return True

def summary_simulations(path=SIMULATION_RESULT_PATH):
    html_str = ""
    # loop all simulations and load their info and add to the html table
    for each_path in os.listdir(path):
        if not check_path_valid(os.path.join(path, each_path)):
            continue
        sim_info = load(os.path.join(path, each_path, 'sim.info'))
        # dataset_with_map = '-' + sim_info['map_info'] if 'map_info' in sim_info and sim_info['map_info'] is not None else ''
        # dataset_with_map = sim_info['dataset'] + dataset_with_map
        dataset_with_map = sim_info['dataset']
        html_str += f"
{sim_info['name']}	" \
                    f"{sim_info['task']}	" \
                    f"{dataset_with_map}	" \
                    f"{sim_info['planner']}	" \
                    f"{sim_info['predictor']}	" \
                    f"{sim_info['status']}	" \
                    f"{sim_info['starting_time']}	"
        html_str += f"{sim_info['ending_time']}	" if sim_info['ending_time'] is not None else "-	"
        # add action drop
        html_str += f'''

                Detail

          '''
        html_str += "
"
    return html_str

def list_simulation_selection(dataset=None, path=SIMULATION_RESULT_PATH):
    if dataset is None:
        return
    html_str = ""
    for each_path in os.listdir(path):
        if not check_path_valid(os.path.join(path, each_path)):
            continue
        if dataset not in each_path:
            continue
        html_str += f"<option value=\"{each_path}\">{each_path}</option>"
    return html_str

def list_scenarios(simulation_name=None):
    if simulation_name is None:
        return
    simulation_path = os.path.join(SIMULATION_RESULT_PATH, simulation_name)
    if not os.path.exists(simulation_path):
        return
    html_str = """

          Scenario id
          Collsion Rate
          Progress (m)

      """
    for each_playback_path in os.listdir(os.path.join(simulation_path,  'playback')):
        loaded_playback = load(os.path.join(simulation_path, 'playback', each_playback_path))
        for each_scenario_id in loaded_playback:
            html_str += "
"
            metric_rst = loaded_playback[each_scenario_id]['metrics']
            task = loaded_playback[each_scenario_id]['info']['task']
            dataset = loaded_playback[each_scenario_id]['info']['dataset']
            collision_rate = len(metric_rst['collided_pairs'])
            progress = metric_rst['progress']
            # jerk = "N/A"  # metric_rst['jerk']
            html_str += f"""
            {each_scenario_id}
            {collision_rate}
            {progress}
            """
            file_name = each_playback_path.split('.playback')[0]
            html_str += f"""

                        Visualize

                """
    html_str += """

"""

    return html_str


import apace as ap
import matplotlib.pyplot as plt
import numpy as np
from fodo import make_fodo

from master_thesis import figure_path

angles = [0, np.pi / 8, np.pi / 4]
fodo = make_fodo(angle=angles[1])
d1, b1, q1, q2 = (fodo[name] for name in ("d1", "b1", "q1", "q2"))
twiss = ap.Twiss(fodo)

steps = 1000
lengths = [0.5, 1.0, 1.5]
n_rows, n_cols = len(lengths) + 1, len(angles)
fig, axs = plt.subplots(
    nrows=n_rows,
    ncols=n_cols,
    figsize=(4 * n_cols, 2.5 * n_rows),
    gridspec_kw={"height_ratios": [0.01, 1, 1, 1]},
)

for ax, angle in zip(axs[0], angles):
    ax.axis("off")
    ax.set_title(
        f"Dipole angle ({np.degrees(angle)}° per cell)", fontweight="bold", pad=0
    )

for column, angle in zip(axs[1:].T, angles):
    b1.angle = angle
    b1.e1 = b1.e2 = 0.5 * angle

    for ax, length in zip(column, lengths):
        d1.length = length
        # breakpoint()
        extent = 0, 2, 0, -2
        results = np.empty((steps, steps))
        for i, q1.k1 in enumerate(np.linspace(*extent[:2], steps)):
            for j, q2.k1 in enumerate(np.linspace(*extent[2:], steps)):
                try:
                    results[i, j] = np.mean(twiss.beta_x) + np.mean(twiss.beta_y)
                except ap.UnstableLatticeError:
                    results[i, j] = np.nan
        image = ax.imshow(
            results.T, extent=extent, origin="lower", vmin=0, vmax=30, cmap="cool"
        )
        ax.set_xlabel(f"$k_\\mathrm{{{q1.name}}}$ / m$^{{-2}}$")
        ax.set_ylabel(f"$k_\\mathrm{{{q2.name}}}$ / m$^{{-2}}$")
        ax.set_title(f"cell length: {fodo.length} m")
        colorbar = fig.colorbar(image, ax=ax)
        colorbar.ax.set_title(r"$\beta_\mathrm{mean}$", fontsize=12, pad=10)

plt.tight_layout()
plt.savefig(figure_path / "necktie-plot.svg")


from numpy import squeeze, real, mean, pi, float16, array, float16, reshape, float32
from scipy import ndimage as ndi
from skimage.filters import gabor_kernel
from skimage.feature import hog
from skimage import feature
import cv2
import numpy as np
from skimage.transform import rescale, resize, downscale_local_mean
import pywt

def hogFeature(normalizedIrisPatch, regions):
    # regions: [(x1, x2), (x3, x4), (x5, x6), ...]
    upperCutHeight = 10
    # HOG Features
    hogFea = []
    for reg in regions:
        croppedImage = normalizedIrisPatch[upperCutHeight:, reg[0]:reg[1]]
        hog_cur = hog(croppedImage, orientations=6, pixels_per_cell=(32, 32), cells_per_block=(1, 1))
        hog_cur = array(hog_cur, float32)
        hogFea.append(hog_cur)

    hogFea = array(hogFea, dtype=float32)
    hogFea = reshape(hogFea, (hogFea.shape[0] * hogFea.shape[1],1))
    hogFea = hogFea.tolist()

    return hogFea

def lbpFeature(normalizedIrisPatch, regions):

    # regions: [(x1, x2), (x3, x4), (x5, x6), ...]
    P = 16
    upperCutHeight = 10
    # LBP Features
    lbpFea = []
    for reg in regions:
        croppedImage = normalizedIrisPatch[upperCutHeight:, reg[0]:reg[1]]
        lbp = feature.local_binary_pattern(croppedImage, 16, 2, method='uniform')
        hist, _ = np.histogram(lbp, normed=True, bins=P + 2, range=(0, P + 2))
        lbpFea.append(hist)

    lbpFea = array(lbpFea, dtype=float32)
    lbpFea = reshape(lbpFea, (lbpFea.shape[0] * lbpFea.shape[1],1))
    lbpFea = lbpFea.tolist()

    return lbpFea

def gaborFeature(normalizedIrisPatch, regions):
    # regions: [(x1, x2), (x3, x4), (x5, x6), ...]
    upperCutHeight = 10
    # Gabor Features
    kernels = []
    freqs = [0.1, 0.2, 0.3, 0.4, 0.5]
    nTheta = 8
    for theta in range(nTheta):
        theta = theta / float16(nTheta) * pi
        sigma = 1
        for frequency in freqs:
            kernel = real(gabor_kernel(frequency, theta=theta, sigma_x=sigma, sigma_y=sigma))
            kernels.append(kernel)

    gaborFea = []

    for reg in regions:
        croppedImage = normalizedIrisPatch[upperCutHeight:, reg[0]:reg[1]]
        gaborFea_cur = []
        for k, kernel in enumerate(kernels):
            filteredIris = ndi.convolve(croppedImage, kernel, mode='wrap')
            gaborFea_cur.append(mean(filteredIris * filteredIris))
        gaborFea_cur = array(gaborFea_cur, float32)
        gaborFea.append(gaborFea_cur)

    gaborFea = array(gaborFea, dtype=float32)
    gaborFea = reshape(gaborFea, (gaborFea.shape[0] * gaborFea.shape[1],1))
    gaborFea =gaborFea.tolist()

    return gaborFea

def extract_image_feature(image, regions, downSampleSize):
    # regions: [(x1, x2), (x3, x4), (x5, x6), ...]
    upperCutHeight = 10
    # Pixel Features
    pixelFea = []
    for reg in regions:
        croppedImage    = image[upperCutHeight:, reg[0]:reg[1]]
        downSampledReg  = rescale(croppedImage, 1.0 / float16(downSampleSize), preserve_range=True)
        pixelFea.append(reshape(downSampledReg, (downSampledReg.shape[0]*downSampledReg.shape[1],)))

    pixelFea = array(pixelFea, dtype=float32)
    pixelFea = reshape(pixelFea, (pixelFea.shape[0]*pixelFea.shape[1], 1))
    pixelFea = pixelFea.tolist()

    return pixelFea


import os
import json
import numpy as np
from ortools.sat.python import cp_model

def solve(all_block, wafer_width, wafer_height):
    ### wafer sampling
    # Number of blocks
    n = len(all_block)

    # wafer Variables
    all_wafer_x_st, all_wafer_y_st, all_wafer_x_ed, all_wafer_y_ed, sampled = [], [], [], [], []
    for i, block in enumerate(all_block):
        if block['x'] == None:
            all_wafer_x_st.append(model.NewIntVar(0, wafer_width - block['w'], f"wafer_x{i}"))
        else:
            all_wafer_x_st.append(model.NewIntVar(block['x'], block['x'], f"wafer_x{i}"))
        all_wafer_x_ed.append(model.NewIntVar(-wafer_width, wafer_width, f"wafer_x_end{i}"))

        if block['y'] == None:
            all_wafer_y_st.append(model.NewIntVar(0, wafer_height - block['h'], f"wafer_y{i}"))
        else:
            all_wafer_y_st.append(model.NewIntVar(block['y'], block['y'], f"wafer_y{i}"))
        all_wafer_y_ed.append(model.NewIntVar(-wafer_height, wafer_height, f"wafer_y_end{i}"))

        if block['x'] == None and block['y'] == None:
            sampled.append(model.NewBoolVar(f"sampled_{i}"))
        else:
            sampled.append(model.NewConstant(1))

    # wafer width & height constraints
    for i, block in enumerate(all_block):
        model.Add(all_wafer_x_ed[i] == all_wafer_x_st[i] + block['w']).OnlyEnforceIf(sampled[i])
        model.Add(all_wafer_y_ed[i] == all_wafer_y_st[i] + block['h']).OnlyEnforceIf(sampled[i])
        model.Add(all_wafer_x_ed[i] <= wafer_width).OnlyEnforceIf(sampled[i])
        model.Add(all_wafer_y_ed[i] <= wafer_height).OnlyEnforceIf(sampled[i])

    # wafer Non-overlapping constraints
    for i in range(n):
        for j in range(i + 1, n):
            wafer_bx_ij = model.NewBoolVar(f"wafer_bx_{i}_{j}")
            wafer_bx_ji = model.NewBoolVar(f"wafer_bx_{j}_{i}")
            wafer_by_ij = model.NewBoolVar(f"wafer_by_{i}_{j}")
            wafer_by_ji = model.NewBoolVar(f"wafer_by_{j}_{i}")

            model.Add(all_wafer_x_ed[i] <= all_wafer_x_st[j] + wafer_width * wafer_bx_ij)
            model.Add(all_wafer_x_ed[j] <= all_wafer_x_st[i] + wafer_width * wafer_bx_ji)
            model.Add(all_wafer_y_ed[i] <= all_wafer_y_st[j] + wafer_height * wafer_by_ij)
            model.Add(all_wafer_y_ed[j] <= all_wafer_y_st[i] + wafer_height * wafer_by_ji)

            model.AddBoolOr([wafer_bx_ij.Not(), wafer_bx_ji.Not(),
                            wafer_by_ij.Not(), wafer_by_ji.Not()])

    ### place to sample panel
    panel_width = 6
    panel_height = 6
    # panel Variables
    all_panel_x_st, all_panel_y_st, all_panel_x_ed, all_panel_y_ed, on_panel = [], [], [], [], []
    for i, block in enumerate(all_block):
        all_panel_x_st.append(model.NewIntVar(0, panel_width - block['w'], f"panel_x{i}"))
        all_panel_x_ed.append(model.NewIntVar(-panel_width, panel_width, f"panel_x_end{i}"))
        all_panel_y_st.append(model.NewIntVar(0, panel_height - block['h'], f"panel_y{i}"))
        all_panel_y_ed.append(model.NewIntVar(-panel_height, panel_height, f"panel_y_end{i}"))
        # on_panel.append(model.NewBoolVar(f"on_panel_{i}"))

    # panel width & height constraints
    for i, block in enumerate(all_block):
        model.Add(all_panel_x_ed[i] == all_panel_x_st[i] + block['w']).OnlyEnforceIf(sampled[i])
        model.Add(all_panel_y_ed[i] == all_panel_y_st[i] + block['h']).OnlyEnforceIf(sampled[i])
        model.Add(all_panel_x_ed[i] <= panel_width).OnlyEnforceIf(sampled[i])
        model.Add(all_panel_y_ed[i] <= panel_height).OnlyEnforceIf(sampled[i])

    # panel Non-overlapping constraints
    for i in range(n):
        for j in range(i + 1, n):
            panel_bx_ij = model.NewBoolVar(f"panel_bx_{i}_{j}")
            panel_bx_ji = model.NewBoolVar(f"panel_bx_{j}_{i}")
            panel_by_ij = model.NewBoolVar(f"panel_by_{i}_{j}")
            panel_by_ji = model.NewBoolVar(f"panel_by_{j}_{i}")

            model.Add(all_panel_x_ed[i] <= all_panel_x_st[j] + panel_width * panel_bx_ij)
            model.Add(all_panel_x_ed[j] <= all_panel_x_st[i] + panel_width * panel_bx_ji)
            model.Add(all_panel_y_ed[i] <= all_panel_y_st[j] + panel_height * panel_by_ij)
            model.Add(all_panel_y_ed[j] <= all_panel_y_st[i] + panel_height * panel_by_ji)

            model.AddBoolOr([panel_bx_ij.Not(), panel_bx_ji.Not(),
                            panel_by_ij.Not(), panel_by_ji.Not()])

    # panel must be filled by blocks
    model.Add(sum(sampled[i] * block['w'] * block['h'] for i, block in enumerate(all_block)) == panel_width * panel_height)

    # Objective function
    wafer_area = wafer_width * wafer_height
    blocks_area = model.NewIntVar(0, wafer_area, "blocks_area")
    model.Add(
        blocks_area == sum(
            sampled[i] *
            block['w'] *
            block['h'] for i, block in enumerate(all_block)))
    num_blocks_sampled = model.NewIntVar(0, n, "num_blocks_sampled")
    model.Add(num_blocks_sampled == sum(sampled[i] for i, block in enumerate(all_block)))

    scale = 1000000
    # wafer_coverage
    wafer_coverage = model.NewIntVar(0, 1 * scale, "wafer_coverage")
    model.AddDivisionEquality(wafer_coverage, blocks_area * scale, wafer_area)

    # block utilization
    block_utilization = model.NewIntVar(0, 1 * scale, "block_utilization")
    model.AddDivisionEquality(
        block_utilization,
        num_blocks_sampled * scale,
        n)

    # model.Maximize(wafer_coverage * (1 / scale) -
    #                block_utilization * (1 / scale))
    # model.Maximize(wafer_coverage * (1 / scale))
    model.Maximize(-block_utilization * (1 / scale))

    # Solve the model
    solver = cp_model.CpSolver()
    status = solver.Solve(model)

    # Print
    if status == cp_model.OPTIMAL:
        wafer_positions = [(solver.Value(all_wafer_x_st[i]), solver.Value(all_wafer_y_st[i]))
                     for i in range(n)]
        panel_positions = [(solver.Value(all_panel_x_st[i]), solver.Value(all_panel_y_st[i]))
                     for i in range(n)]
        print(f"wafer_positions: {wafer_positions}\n"
              f"panel_positions: {panel_positions}\n"
              f"num_blocks_sampled: {solver.Value(num_blocks_sampled)}\n"
              f"sampled: {[solver.Value(sampled[i]) for i in range(n)]}\n"
              f"wafer_coverage: {solver.Value(wafer_coverage) / scale}\n"
              f"block_utilization: {solver.Value(block_utilization) / scale}\n"
              f"objective: {solver.ObjectiveValue()}")
        all_block_sampled = []
        for i, block in enumerate(all_block):
            if not solver.Value(sampled[i]):
                continue
            block['x'] = solver.Value(all_wafer_x_st[i])
            block['y'] = solver.Value(all_wafer_y_st[i])
            all_block_sampled.append(block)
        result = {}
        result['width'] = wafer_width
        result['height'] = wafer_height
        result["block"] = all_block_sampled
        with open(os.path.join(result_path, file_name), 'w') as fp:
            json.dump(result, fp, indent=4)
    elif cp_model.INFEASIBLE:
        print("INFEASIBLE")

if __name__ == "__main__":
    data_path = "block_data"
    result_path = "result"
    # file_name = "0004.json"
    file_name = "0005.json"
    # file_name = "0001_d=10.json"
    # Create the model
    model = cp_model.CpModel()
    with open(os.path.join(data_path, file_name), 'r') as fp:
        data = json.load(fp)
    wafer_width = data["width"]
    wafer_height = data["height"]
    all_block = data["block"]
    solve(all_block, wafer_width, wafer_height)


#! /bin/python

import os
import sys
import json
import numpy as np
import luigi

import vigra
import nifty
import nifty.tools as nt

import cluster_tools.utils.volume_utils as vu
import cluster_tools.utils.function_utils as fu
from cluster_tools.cluster_tasks import SlurmTask, LocalTask, LSFTask

#
# Orphan Filter Tasks
#

class OrphanAssignmentsBase(luigi.Task):
    """ OrphanAssignments base class
    """

    task_name = 'orphan_assignments'
    src_file = os.path.abspath(__file__)
    allow_retry = False

    graph_path = luigi.Parameter()
    graph_key = luigi.Parameter()
    assignment_path = luigi.Parameter()
    assignment_key = luigi.Parameter()
    output_path = luigi.Parameter()
    output_key = luigi.Parameter()
    relabel = luigi.BoolParameter(default=False)
    #
    dependency = luigi.TaskParameter()

    def requires(self):
        return self.dependency

    def run_impl(self):
        # get the global config and init configs
        shebang = self.global_config_values()[0]
        self.init(shebang)

        # load the task config
        config = self.get_task_config()

        # update the config with input and graph paths and keys
        # as well as block shape
        config.update({'assignment_path': self.assignment_path,
                       'assignment_key': self.assignment_key,
                       'graph_path': self.graph_path,
                       'graph_key': self.graph_key,
                       'output_path': self.output_path,
                       'output_key': self.output_key,
                       'relabel': self.relabel})

        n_jobs = 1
        # prime and run the jobs
        self.prepare_jobs(n_jobs, None, config)
        self.submit_jobs(n_jobs)

        # wait till jobs finish and check for job success
        self.wait_for_jobs()
        self.check_jobs(n_jobs)

class OrphanAssignmentsLocal(OrphanAssignmentsBase, LocalTask):
    """ OrphanAssignments on local machine
    """
    pass

class OrphanAssignmentsSlurm(OrphanAssignmentsBase, SlurmTask):
    """ OrphanAssignments on slurm cluster
    """
    pass

class OrphanAssignmentsLSF(OrphanAssignmentsBase, LSFTask):
    """ OrphanAssignments on lsf cluster
    """
    pass

#
# Implementation
#

def orphan_assignments(job_id, config_path):

    fu.log("start processing job %i" % job_id)
    fu.log("reading config from %s" % config_path)

    # get the config
    with open(config_path) as f:
        config = json.load(f)

    # load from config
    assignment_path = config['assignment_path']
    assignment_key = config['assignment_key']
    graph_path = config['graph_path']
    graph_key = config['graph_key']
    output_path = config['output_path']
    output_key = config['output_key']
    relabel = config['relabel']
    n_threads = config.get('threads_per_job', 1)

    # load the uv-ids and assignments
    with vu.file_reader(graph_path) as f:
        ds = f['%s/edges' % graph_key]
        ds.n_threads = n_threads
        uv_ids = ds[:]
    with vu.file_reader(assignment_path) as f:
        ds = f[assignment_key]
        ds.n_threads = n_threads
        chunks = ds.chunks
        assignments = ds[:]

    n_new_nodes = int(assignments.max()) + 1
    # find the new uv-ids
    edge_mapping = nt.EdgeMapping(uv_ids, assignments, numberOfThreads=n_threads)
    new_uv_ids = edge_mapping.newUvIds()

    # find all orphans = segments that have node degree one
    ids, node_degrees = np.unique(new_uv_ids, return_counts=True)
    orphans = ids[node_degrees == 1]
    n_orphans = len(orphans)
    fu.log("Found %i orphans of %i clusters" % (n_orphans, n_new_nodes))

    # make graph for fast neighbor search
    graph = nifty.graph.undirectedGraph(n_new_nodes)
    graph.insertEdges(new_uv_ids)

    orphan_assignments = np.array([next(graph.nodeAdjacency(orphan_id))[0]
                                   for orphan_id in orphans],)
    assert len(orphan_assignments) == n_orphans, "%i, %i" % (len(orphan_assignments), n_orphans)
    assignments[orphans] = orphan_assignments.astype('uint64')

    if relabel:
        vigra.analysis.relabelConsecutive(assignments, out=assignments,
                                          start_label=1, keep_zeros=True)

    with vu.file_reader(output_path) as f:
        ds = f.require_dataset(output_key, shape=assignments.shape, chunks=chunks,
                               compression='gzip', dtype='uint64')
        ds[:] = assignments

    fu.log_job_success(job_id)

if __name__ == '__main__':
    path = sys.argv[1]
    assert os.path.exists(path), path
    job_id = int(os.path.split(path)[1].split('.')[0].split('_')[-1])
    orphan_assignments(job_id, path)


import re
from validate_email import validate_email

TAG_RE = re.compile(r'<[^>]+>')
def remove_tags(text):
    return TAG_RE.sub('', text)

class Email():
  EMAIL_FIELDS = ["to", "from"]
  FIELDS = "to to_name from from_name subject body".split()
  def __init__(self, raw_data):
    self.populate_fields(raw_data)
    self.validate_emails()
    self.sanitize_body()

  def populate_fields(self, raw_data):
    fields = "to to_name from from_name subject body".split()
    for key in self.FIELDS:
      if raw_data.has_key(key) and isinstance(raw_data[key], basestring):
        setattr(self, key, str(raw_data[key]))
      else:
        raise Exception("Error, invalid data for '{}'.".format(key))

  def validate_emails(self):
    for field in self.EMAIL_FIELDS:
      address = getattr(self, field)
      if not validate_email(address):
        raise Exception("Error, invalid email '{}'".format(address))

  def sanitize_body(self):
    self.body = remove_tags(self.body)


#!/usr/bin/env python
# coding: utf-8

# In[9]:

import json
with open('file1.json','r') as a:
    data1 = a.read()
obj1 = json.loads(data1)
with open('file2.json','r') as a:
    data2 = a.read()
obj2 = json.loads(data2)
dlt = {i: obj1[i] for i in obj1 if i in obj2 and obj1[i] != obj2[i]}
if len(dlt):
    print ("Есть различие!\nJSON 1 | JSON 2")
    for key, value in dlt.items():
        print (key, "->", value, '|',key, "->", obj2[key])

# In[ ]:


from popsycle import synthetic
import numpy as np
import matplotlib.pyplot as plt
from astropy.table import Table
import h5py

def test_h5_output(extra_col= True):
    """"
    Parameters
    ----------

     extra_col : boolean, defaults to False
        Tells the code whether or not the new h5 file will have additional columns (ie does the new version of
        popsycle give more information than before
    """

    #find the test files
    test_data_dir = '/u/samrose/scratch/test_files/'
    ebf_file= test_data_dir + 'h5_reference.ebf'
    reference_h5_file= test_data_dir + 'h5_reference.h5'

    #create the new h5 file by running popsycle
    synthetic.perform_pop_syn(ebf_file = ebf_file,
                              output_root = 'test',
                              iso_dir = '/u/casey/scratch/work/microlens/popsycle_test/isochrones/',
                              bin_edges_number = None, overwrite = True, seed=42);

    #read in the data from the reference h5 file
    hfr = h5py.File(reference_h5_file, 'r')
    ref_dset = np.concatenate((hfr['l0b0'], hfr['l0b1'], hfr['l1b0'], hfr['l1b1']),
                              axis=1)
    hfr.close()

    #read in the data from the test h5 file created by popsycle
    hft = h5py.File('test.h5', 'r')
    test_dset = np.concatenate((hft['l0b0'], hft['l0b1'], hft['l1b0'], hft['l1b1']),
                              axis=1)
    hft.close()

    #see if we have the right number of columns
    if test_dset.shape[0] != ref_dset.shape[0] and not extra_col:
            assert test_dset.shape[0] == ref_dset.shape[0], "the h5 files are not the same size. Run again with extra_col=True if you have added columns)"

    #test to see whether the files are the same
    matched_col=0 #initialize matched_col counter
    for i in range(0, ref_dset.shape[0]):
        test_col = test_dset[i,:]
        ref_col = ref_dset[i, :]
        if test_col.all() == ref_col.all():
            matched_col = matched_col+1

        #check to see if disagreements are because of nans
        else:
            bad_idxs = np.where(ref_col != test_col)
            ref_nan_idx = np.where(ref_col == np.nan)
            test_nan_idx = np.where(test_col == np.nan)
            if test_nan_idx.all() == ref_nan_idx.all() and bad_idxs.all() == ref_nan_idx.all():
                matched_col = matched_col+1
            else:
                matched_col= matched_col
                assert test_nan_idx.all() == ref_nan_idx.all(), "Files do not have nan values at the same indices"
                assert bad_idxs.all() == ref_nan_idx.all(), "Coulumns disagree at non-nan values"

    assert matched_col == ref_dset.shape[0], "The new test h5 file does not match the reference file!"

    return


import logging
import threading
import types
from collections import namedtuple
from hashlib import sha256
from time import sleep, time

from goTenna.payload import BinaryPayload, CustomPayload
from termcolor import colored

import config
from utilities import de_segment, naturalsize

logger = logging.getLogger("MSGS")
mesh_logger = logging.getLogger("MESH")

def handle_message(conn, queue):
    """
    Handle messages received over the mesh network
    :param conn: the lntenna.gotenna.Connection instance
    :param queue: a queue.Queue() containing messages
    :return: result of message handling
    """
    while True:
        if queue.empty():
            sleep(0.15)
        else:
            message = queue.get().message

            if isinstance(message.payload, CustomPayload):
                print(message)
            elif isinstance(message.payload, BinaryPayload):
                payload = message.payload._binary_data
                digest = sha256(payload).hexdigest()
                conn.bytes_received += len(payload)
                if config.DEBUG:
                    mesh_logger.info(
                        colored(
                            f"Received {naturalsize(len(payload))} - {digest}", "cyan"
                        )
                    )
                else:
                    mesh_logger.info(
                        colored(f"Received {naturalsize(len(payload))}", "cyan")
                    )
                if not payload[0:4] in config.VALID_MSGS:
                    logger.error(
                        "Message magic not found in VALID_MSGS. Discarding message"
                    )
                    return
                conn.events.send_via_socket.put(payload[4:])
            else:
                payload = message.payload.message
                # test for jumbo:
                jumbo = True if payload.startswith("sm/") else False
                if jumbo:
                    handle_jumbo_message(conn, message)
                    return
                else:
                    logger.error("Unhandled payload type received:")
                    logger.error(payload)

def handle_jumbo_message(conn, message):
    """Handle a jumbo message received.
    """
    payload = message.payload.message
    # TODO: this cuts out all sender and receiver info -- ADD SENDER GID
    logger.info(f"Received jumbo message fragment")
    prefix, seq, length, msg = payload.split("/")

    # if a jumbo monitor thread is not running, start one
    if conn.jumbo_thread.is_alive():
        pass
    else:
        conn.events.jumbo_len = length
        conn.jumbo_thread = None
        conn.jumbo_thread = threading.Thread(
            target=monitor_jumbo_msgs, daemon=True, args=[conn]
        )
        conn.jumbo_thread.start()
    # add the message to the events.jumbo queue
    conn.events.jumbo.append(payload)
    return

def monitor_jumbo_msgs(conn, timeout=210):
    logger.debug("Starting jumbo message monitor thread")
    start = time()
    missing = True
    while True and time() < start + timeout:
        # logger.info(
        #     f"received: {len(conn.events.jumbo)} of {conn.events.jumbo_len} "
        #     f"jumbo messages"
        # )
        if (
            len(conn.events.jumbo) == int(conn.events.jumbo_len)
            and len(conn.events.jumbo) is not 0
        ):
            missing = False
            # give handle_message the attributes it expects
            jumbo_message = types.SimpleNamespace()
            jumbo_message.payload = types.SimpleNamespace()
            # reconstruct the jumbo message
            jumbo_message.payload.message = de_segment(conn.events.jumbo)
            # send it back through handle_message
            logger.info(f"Jumbo message payload reconstituted")
            handle_message(conn, jumbo_message)
            break
        sleep(0.2)
    # reset jumbo events after timeout
    conn.events.init_jumbo()
    if missing:
        logger.error(
            "Did not receive all jumbo messages require for re-assembly. "
            "Please request the message again from the remote host."
        )
    return

"""
Message Structure:

Size    | Description
-----------------------
4       | Magic / Protocol
16      | Host
2       | Port
4       | Checksum / Peer (ID)

This will associate this checksum (peer) with this ip address/port configuration, for
this protocol.

Future messages must all be prefixed with `Checksum`.

Messages not prefixed with a valid Magic or Checksum will be discarded.
"""

# checksums = {}
# Peer = namedtuple("Peer", ["host", "port", "protocol"])
#
#
# def handle_binary_msg(msg):
#     # throw away the message if it's not in magic or the checksum DB
#     prefix = msg[0:4]
#     if prefix not in MAGIC and checksums:
#         print(f"Message prefix unknown: {msg[0:4]}")
#         return
#
#     if prefix in MAGIC:
#         if not len(msg) == 26:
#             print(f"Invalid message length for magic negotiation: {len(msg)}")
#             return
#         # add the host, port, protocol to the peer's entry in checksums
#         checksums[prefix] = Peer(msg[4:20], msg[20:22], msg[0:4])
#         print(f"Peer {prefix} added to in-memory peer dictionary")
#
#     elif prefix in checksums:
#         # if ltng protocol, just strip the header and return it for now
#         if checksums[prefix] == b"ltng":
#             print(f"Peer {prefix}'s message stripped and returned")
#             return msg[4:]


import time

from osv import osv
from osv import fields
from tools import config
from tools.translate import _
from datetime import datetime
from datetime import timedelta

class hr_payroll_declar(osv.osv):
    '''
    Decleration Form
    '''
    _name = 'hr.payroll.declare'
    _description = 'Decleration Form'

    _columns = {
        'name':fields.char('Name', size=1024, required=False),
        'company_id':fields.many2one('res.company', 'Company', required=True),
        'employee_id':fields.many2one('hr.employee', 'Employee', required=True),
        'income_sal': fields.float('Income by Salary', digits=(16, int(config['price_accuracy'])), readonly=True),
        'income_ids':fields.one2many('hr.payroll.declare.line', 'income_id', 'Source Of Income', required=False),
        'investment_ids':fields.one2many('hr.payroll.declare.line', 'invest_id', 'Investments', required=False),
        'claim_ids':fields.one2many('hr.payroll.declare.line', 'claim_id', 'Allowance to Claime', required=False),
        'date': fields.date('Date'),

        'income': fields.float('Taxable Income', digits=(16, int(config['price_accuracy'])), readonly=True),
        'investment': fields.float('Total Investment', digits=(16, int(config['price_accuracy'])), readonly=True),
        'claims': fields.float('Total Allowance Claims', digits=(16, int(config['price_accuracy'])), readonly=True),

        'state':fields.selection([
            ('draft','Draft'),
            ('pending','Waiting for Review'),
            ('pending','Approved by HR'),
            ('done','Confirm'),
        ],'State', select=True, readonly=True),
        'note': fields.text('Description'),
    }

    def get_basic(self, cr, uid, ids, context):
        res = {}
        for rs in self.browse(cr, uid, ids, context):
            period_id = self.pool.get('account.period').search(cr,uid,[('date_start','<=',time.strftime('%Y-%m-%d')),('date_stop','>=',time.strftime('%Y-%m-%d'))])[0]
            fiscalyear_id = self.pool.get('account.period').browse(cr, uid, period_id).fiscalyear_id
            sql_req= '''
                SELECT c.id as id, c.wage as wage, function as function, c.date_start as start, c.date_end as end
                FROM hr_contract c
                  LEFT JOIN hr_employee emp on (c.employee_id=emp.id)
                  LEFT JOIN hr_contract_wage_type cwt on (cwt.id = c.wage_type_id)
                  LEFT JOIN hr_contract_wage_type_period p on (cwt.period_id = p.id)
                WHERE
                  (emp.id=%s) AND
                  (date_start >= %s) AND
                  (date_end IS NULL OR date_end <= %s)
                '''
            cr.execute(sql_req, (rs.employee_id.id, fiscalyear_id.date_start, fiscalyear_id.date_stop))
            contracts = cr.dictfetchall()
            if not contracts:
                raise osv.except_osv(_('Contract Error !'), _('No Contract Defined for : %s ' % (rs.employee_id.name)))
            total = 0.0

            line_ids = []
            for lines in rs.claim_ids:
                line_ids += [lines.head_id.id]

            for ct in contracts:
                allow = 0.0

                d1 = ct['start']
                d2 = ct['end'] or fiscalyear_id.date_stop
                td = datetime.fromtimestamp(time.mktime(time.strptime(d2, '%Y-%m-%d'))) - datetime.fromtimestamp(time.mktime(time.strptime(d1, '%Y-%m-%d')))
                total += (td.days / 30) * ct['wage']

#                ct = self.pool.get('hr.contract').browse(cr, uid, ct['id'])
#                for line in ct.function.line_ids:
#                    if line.category_id.id in line_ids:
#                        if line.amount_type == 'fix':
#                            allow += (td.days / 30) * line.amount
#                        elif line.amount_type == 'per':
#                            allow += (total * line.amount)
#                        print 'XXXXXXXXXXXXXXXXXXXXXXX : ', line.name, allow

            res[rs.id] = total
        return res

    def write(self, cr, user, ids, vals, context=None):
        res = self.get_basic(cr, user, ids, context)
        for id in ids:
            vals['income_sal'] = res[id]
            super(hr_payroll_declar, self).write(cr, user, [id], vals, context)

        return res

hr_payroll_declar()

class hr_payroll_declare_line(osv.osv):
    '''
    Decleration Line
    '''
    _name = 'hr.payroll.declare.line'
    _description = 'Decleration Line'

    def _function_call(self, cr, uid, ids, field_names, arg, context={}):
        res = {}
        for rs in self.browse(cr, uid, ids, context):
            val = 0.0

            if rs.income_id:
                pass
            elif rs.invest_id:
                pass
            elif rs.claim_id:
                if rs.head_id.calc_type == 'min_max':
                    if rs.amount < rs.head_id.min:
                        val = rs.head_id.min
                    elif rs.amount >= rs.head_id.min and rs.amount <= rs.head_id.max:
                        val = rs.amount
                    elif rs.amount > rs.head_id.max:
                        val = rs.head_id.max

            res[rs.id] = val

        return res

    _columns = {
        'name':fields.char('Name', size=64, required=False),
        'note': fields.text('Description'),
        'income_id':fields.many2one('hr.payroll.declare', 'Income', required=False),
        'invest_id':fields.many2one('hr.payroll.declare', 'Investment', required=False),
        'claim_id':fields.many2one('hr.payroll.declare', 'Allowance Claims', required=False),
        'amount': fields.float('Amount', digits=(16, int(config['price_accuracy']))),
        'allow': fields.float('Allowence', digits=(16, int(config['price_accuracy']))),
        'allow_amount': fields.function(_function_call, method=True, type='float', digits=(16, int(config['price_accuracy'])), string='Allow Amount'),
        'head_id':fields.many2one('hr.allounce.deduction.categoty', 'Allowance / Deduction', required=True),

    }
hr_payroll_declare_line()

class payment_category(osv.osv):
    '''
    Allowance Deduction Categoty
    '''
    _inherit = 'hr.allounce.deduction.categoty'

    _columns = {
        'calc_type':fields.selection([
            ('min_max','Min / Max'),
            ('stmt','List of Calculations'),
            ('range','Selection from Range'),
        ],'Calculation Type', select=True, readonly=False),
        'min': fields.float('Min Value', digits=(16, int(config['price_accuracy']))),
        'max': fields.float('Max Value', digits=(16, int(config['price_accuracy']))),

        'stmt_ids':fields.one2many('hr.payroll.declare.stmt', 'category_id', 'Functions', required=False),
        'stmt_select':fields.selection([
            ('min','Minimum'),
            ('max','Maximum'),
            ('avg','Average'),
        ],'Selection Method', select=True, readonly=False),
    }
    _defaults = {
        'stmt_select': lambda *a: 'min',
        'calc_type': lambda *a: 'min_max'
    }
payment_category()

class payment_stmt(osv.osv):
    '''
    Open ERP Model
    '''
    _name = 'hr.payroll.declare.stmt'
    _description = 'Payroll Calculations'

    _columns = {
        'category_id':fields.many2one('hr.allounce.deduction.categoty', 'Category', required=True),
        'name':fields.char('Expression', size=1024, required=True, readonly=False),
        'sequence': fields.integer('Sequence'),
        'active':fields.boolean('Active', required=False),
    }
    _defaults = {
        'sequence': lambda *a: 5,
        'active': lambda *a: True
    }
payment_stmt()


import torch
import torchvision
import gym

import random
import torch.nn as nn
import torch
from torch.autograd import Variable
import torch.autograd as autograd
import torch.nn.functional as F
import gym
import random
import heapq

from gym.envs.registration import register

register(
    id='FrozenLakeNotSlippery-v0',
    entry_point='gym.envs.toy_text:FrozenLakeEnv',
    kwargs={'map_name': '4x4', 'is_slippery': False},
    max_episode_steps=100,
    reward_threshold=0.78,  # optimum = .8196
)

# env = gym.make('FrozenLake8x8-v0')
# env = gym.make('FrozenLake-v0')
env = gym.make('FrozenLakeNotSlippery-v0')
env.render()

use_cuda = torch.cuda.is_available()
FloatTensor = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor
LongTensor = torch.cuda.LongTensor if use_cuda else torch.LongTensor
ByteTensor = torch.cuda.ByteTensor if use_cuda else torch.ByteTensor
Tensor = FloatTensor

class value_net(nn.Module):
    def __init__(self):
        super(value_net, self).__init__()
        bias_on = True
        self.linear1 = nn.Linear(16, 20, bias=bias_on)
        self.linear2 = nn.Linear(20, 40, bias=bias_on)
        self.linear3 = nn.Linear(40, 1, bias=bias_on)
        # self.dropout = nn.Dropout(p=0.5)

    def forward(self, x):
        # --- 0000 ---- 0000 >>>  z-score normalization
        x = self.linear1(x)

        x_avg = torch.sum(x) / 20
        x_minus_x_avg = x - x_avg
        x_std = torch.sum(torch.pow(x_minus_x_avg, 2)) / 20
        epsilon = 0.0000001
        x_norm = (x_minus_x_avg) / (torch.sqrt(x_std) + epsilon)
        x = torch.tanh(x_norm)

        x = self.linear2(x)

        # x_avg = torch.sum(x) / 40
        # x_minus_x_avg = x - x_avg
        # x_std = torch.sum(torch.pow(x_minus_x_avg, 2)) / 40
        # x_norm = (x_minus_x_avg) / (torch.sqrt(x_std) + epsilon)
        x = torch.tanh(x)
        x = self.linear3(x)

        return x.view(-1, 1)

class policy_net(nn.Module):
    def __init__(self):
        super(policy_net, self).__init__()
        bias_on = True
        self.linear1 = nn.Linear(16, 20, bias=bias_on)
        self.linear2 = nn.Linear(20, 40, bias=bias_on)
        self.linear3 = nn.Linear(40, 4, bias=bias_on)
        # self.dropout = nn.Dropout(p=0.5)

    def forward(self, x):
        # --- 0000 ---- 0000 >>>  z-score normalization
        x = self.linear1(x)

        x_avg = torch.sum(x) / 20
        x_minus_x_avg = x - x_avg
        x_std = torch.sum(torch.pow(x_minus_x_avg, 2)) / 20
        epsilon = 0.0000001
        x_norm = (x_minus_x_avg) / (torch.sqrt(x_std) + epsilon)
        x = torch.tanh(x_norm)

        x = self.linear2(x)

        # x_avg = torch.sum(x) / 40
        # x_minus_x_avg = x - x_avg
        # x_std = torch.sum(torch.pow(x_minus_x_avg, 2)) / 40
        # x_norm = (x_minus_x_avg) / (torch.sqrt(x_std) + epsilon)
        x = torch.tanh(x)
        x = self.linear3(x)
        return x.view(-1, 4)

from collections import namedtuple

Transition = namedtuple('Transition',
                        ('state', 'action', 'log_prob', 'action_prob', 'log_action_prob', 'next_state', 'reward',
                         'entropy_impact', 'done'))

class ReplayMemory(object):

    def __init__(self, capacity):
        self.capacity = capacity
        self.memory = []
        self.position = 0

    def push(self, *args):
        """Saves a transition."""
        if len(self.memory) < self.capacity:
            self.memory.append(None)
        self.memory[self.position] = Transition(*args)
        self.position = (self.position + 1) % self.capacity

    def sample(self, batch_size):
        return random.sample(self.memory, batch_size)

    def __len__(self):
        return len(self.memory)

class ReplayMemoryNoReplacement(object):

    def __init__(self, capacity):
        self.h = []

    def push(self, *args):
        random_index = random.random()
        heapq.heappush(self.h, (random_index, Transition(*args)))

    def sample(self, batch_size):
        result = []
        for i in range(batch_size):
            result.append(heapq.heappop(self.h)[1])

        return result

    def __len__(self):
        return len(self.h)

class ReplayMemoryNew(object):
    def __init__(self, capacity):
        self.h = []
        self.capacity = capacity

    def push(self, *args):
        tran = Transition(*args)
        self.push_transition(tran)

    def push_transition(self, tran):
        if self.capacity <= len(self.h):
            heapq.heappop(self.h)
        random_index = random.random()
        heapq.heappush(self.h, (random_index, tran))

    def sample(self, batch_size):
        result = []
        for i in range(batch_size):
            el = heapq.heappop(self.h)[1]
            result.append(el)
            heapq.heappush(self.h, (random.random(), el))
        return result

    def __len__(self):
        return len(self.h)

def print_v_table():
    for i in range(16):
        # st = np.array(get_state_repr(i))
        # st = np.expand_dims(st, axis=0)
        st = get_state_repr(i)
        v_net.eval()
        action_probs = v_net(FloatTensor(st))
        # action_probs = F.softmax(action_probs, dim=1)
        outp = " state (" + str(i) + ") "
        n = 0
        for tensr in action_probs:
            for cell in tensr:
                outp = outp + " A[" + str(n) + "]:(" + str(cell.item()) + ")"
                n += 1
        print(outp)

def print_pi_table():
    for i in range(16):
        # st = np.array(get_state_repr(i))
        # st = np.expand_dims(st, axis=0)
        st = get_state_repr(i)
        pi_net.eval()
        action_probs = pi_net(FloatTensor(st))
        action_probs = F.softmax(action_probs, dim=1)
        outp = " state (" + str(i) + ") "
        n = 0
        for tensr in action_probs:
            for cell in tensr:
                outp = outp + " A[" + str(n) + "]:(" + str(cell.item()) + ")"
                n += 1
        print(outp)

# def get_state_repr(state_idx):
#     return state_idx * 13

import gym
import numpy as np
import torch.optim as optim
from torch.distributions import Categorical
import random

random.seed(1999)
import math
import torch
from torch.optim.lr_scheduler import StepLR

# custom weights initialization
def weights_init(m):
    classname = m.__class__.__name__
    # print classname
    # print q_net
    if classname.find('Linear') != -1:
        m.weight.data.normal_(0.0, 0.02)
        if not m.bias is None:
            m.bias.data.normal_(0.0, 0.02)

def get_state_repr(state_idx):
    state = np.zeros(16)
    state[state_idx] = 1
    return state

def get_index_repr(state):
    return np.argwhere(state==1).item()

# if gpu is to be used
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

BATCH_SIZE = 300
GAMMA = 0.99
TARGET_UPDATE = 1000
PRINT_OUT_TIMES = 1000
ENTROPY_REDUCTION_STEPS = 100000.0
NUM_EPISODES = 10000000
# NUM_STEPS_VALUE_FUNCTION_LEARNS = NUM_EPISODES
#NUM_STEPS_VALUE_FUNCTION_LEARNS = (ENTROPY_REDUCTION_STEPS * 1)
NUM_STEPS_VALUE_FUNCTION_LEARNS = 1

v_net = value_net()
v_net.apply(weights_init)
v_net.to(device)
target_v_net = value_net()
target_v_net.load_state_dict(v_net.state_dict())
target_v_net.to(device)
pi_net = policy_net()
pi_net.apply(weights_init).to(device)

# prepare for optimizer, merge both networks parameters

# parameters = set()
# for net_ in [v_net, pi_net]:
#     parameters |= set(net_.parameters())

# optimizer = optim.RMSprop(online_net.parameters(), lr=0.001)

# optimizer = optim.Adam(parameters, lr=0.0001)

v_optimizer = optim.Adam(v_net.parameters(), lr=0.0001)
pi_optimizer = optim.Adam(pi_net.parameters(), lr=0.00001)

# scheduler = StepLR(v_optimizer, step_size=10000, gamma=0.5)

MEMORY_SIZE = 2000
# memory = ReplayMemoryNoReplacement(MEMORY_SIZE)
memory = ReplayMemoryNew(MEMORY_SIZE)
# memory = ReplayMemory(MEMORY_SIZE)

value_loss_cum = []

def optimize(k):
    if len(memory) < BATCH_SIZE:
        return

    transitions = memory.sample(BATCH_SIZE)
    # Transpose the batch (see http://stackoverflow.com/a/19343/3343043 for
    # detailed explanation).
    batch = Transition(*zip(*transitions))

    # Compute a mask of non-final states and concatenate the batch elements
    # final_mask = torch.tensor(tuple(map(lambda d: d is True,batch.done)), device=device, dtype=torch.bool).unsqueeze(1)
    # final_mask_list = [d for d in batch.done if d is True]
    final_mask = torch.tensor(tuple(map(lambda d: get_index_repr(d) in [5,7,11,12], batch.next_state)), device=device, dtype=torch.bool).unsqueeze(1)
    final_mask_list = [d for d in batch.done if get_index_repr(d) in [5,7,11,12]]
    # Compute states that are final.
    #     next_state_final_mask = torch.tensor(tuple(map(lambda d: (d) in [5,7,11,12,15],
    #                                           batch.next_state)), device=device, dtype=torch.uint8).unsqueeze(1)
    #     next_state_finak_list = [d for d in batch.next_state if d in [5,7,11,12,15] ]

    # Unpack the parameters from the memory

    state_batch = FloatTensor(batch.state)
    state_batch = state_batch.view(BATCH_SIZE, 16)
    next_state_batch = FloatTensor(batch.next_state)
    next_state_batch = next_state_batch.view(BATCH_SIZE, 16)
    action_batch = LongTensor(batch.action).view(BATCH_SIZE, 1)
    reward_batch = Tensor(batch.reward).view(BATCH_SIZE, 1)
    entropy_impact_batch = FloatTensor(batch.entropy_impact).view(BATCH_SIZE, 1)
    # log_prob_batch = torch.cat(batch.log_prob).view(BATCH_SIZE, 1)
    # action_probs_batch = torch.cat(batch.action_prob).view(BATCH_SIZE,4)
    # log_action_probs_batch = torch.cat(batch.log_action_prob).view(BATCH_SIZE,4)

    # FIRST , calculate V(next_state)and backpropagate MSE on V

    target_v_net.eval()
    v_next = target_v_net(next_state_batch).detach()
    # v_next[next_state_final_mask] = torch.zeros(len(next_state_finak_list), device=device).view(len(next_state_finak_list))
    v_next[final_mask] = torch.zeros(len(final_mask_list), device=device).view(len(final_mask_list))

    ##HACK FIXING expected value
    #     v_current_fixed = [get_expected_value_fixed(_st) for _st in batch.state]
    #     v_current_fixed = FloatTensor(v_current_fixed).view(BATCH_SIZE,1)
    ##HACK FIXING expected value

    ##HACK FIXING current value
    #     v_next_fixed = [get_expected_value_fixed(_st) for _st in batch.next_state]
    #     v_next_fixed = FloatTensor(v_next_fixed).view(BATCH_SIZE,1)
    # v_next = v_next_fixed
    ##HACK FIXING current value

    expected_value = reward_batch + v_next * GAMMA

    ##HACK FIXING expected value
    # expected_value = expected_value_fixed
    ##HACK FIXING expected value

    # calculate V(current_state)
    #if k <= NUM_STEPS_VALUE_FUNCTION_LEARNS:
    #    v_net.train()
    #else:
    #    v_net.eval()
    v_net.train()

    v_current = v_net(state_batch)

    # backpropagate:
    value_loss = torch.sum((expected_value - v_current) ** 2)

    v_optimizer.zero_grad()
    value_loss.backward()  # keep graph for policy net optimizer
    v_optimizer.step()

    # if k <= NUM_STEPS_VALUE_FUNCTION_LEARNS:
    #     v_optimizer.zero_grad()
    #     # value_loss.backward(retain_graph=True) # keep graph for policy net optimizer
    #     value_loss.backward()  # keep graph for policy net optimizer
    #     v_optimizer.step()
        # scheduler.step()

    value_loss_cum.append(value_loss.item())

    v_current = v_current.detach()

    ##HACK FIXING expected value
    # v_current = v_current_fixed
    ##HACK FIXING expected value

    # SECOND, calculate gradient loss:
    # H(X) = P(X) log ( P(X) )

    # calculate the action probability
    actions_distr = pi_net(state_batch)
    actions_prob_batch = torch.softmax(actions_distr, dim=1)
    log_actions_prob_batch = torch.log_softmax(actions_distr, dim=1)

    action_batch = action_batch
    action_mask = FloatTensor(BATCH_SIZE, 4).zero_()
    action_mask.scatter_(1, action_batch, 1)  # This will have shape (BATCH_SIZE, 4), and its contents will be
    # like : [[0,0,1,0],[1,0,0,0],...]
    # log_prob_batch = log_actions_prob_batch.gather(1,action_batch)
    log_prob_batch = torch.sum(log_actions_prob_batch * action_mask, dim=1).view(BATCH_SIZE,
                                                                                1)  # sum up across rows (ending tensor is shape (BATCH_SIZE, 1))

    entropy = entropy_impact_batch * torch.sum(actions_prob_batch * log_actions_prob_batch)

    #policy_loss = torch.sum(-log_prob_batch * (expected_value - v_current) + entropy)
    policy_loss = torch.sum(-log_prob_batch * (expected_value - v_current))

    pi_optimizer.zero_grad()
    policy_loss.backward()
    pi_optimizer.step()

    return policy_loss.item(), value_loss.item()

score = []
times_trained = 0
times_reach_goal = 0
steps_done = 0

policy_loss_avg = [1.0]
v_loss_avg = [1.0]

TARGET_UPDATE = 1000

for k in range(NUM_EPISODES):
    done = False
    observation = env.reset()
    # observation, reward, done, info = env.step(env.action_space.sample()) # take a random action
    reward = 0
    episode_step = 0
    # print("b")
    I = 1.0
    # entropy_impact = (ENTROPY_REDUCTION_STEPS - k) / ENTROPY_REDUCTION_STEPS
    if k == 0:
        entropy_impact = 1.0
    else:
        entropy_impact = min(1, (1 / (k * 0.005)))

    if k > ENTROPY_REDUCTION_STEPS:
        entropy_impact = 0.0

    # test entropy always 0
    # entropy_impact = 0.0

    # entropy_impact = 0.0
    # if entropy_impact < 0.0:
    #    entropy_impact = 0
    while not done:
        # print("c")
        steps_done += 1

        # Get action from pi
        # np_observation = np.array(get_state_repr(observation))
        # np_observation = np.expand_dims(np_observation, axis=0)
        np_observation = get_state_repr(observation)
        # print(np_observation)
        observation_tensor = FloatTensor(np_observation)

        # action distribution
        pi_net.eval()
        action_distr = pi_net(observation_tensor)
        action_probs = torch.softmax(action_distr, dim=1)
        log_action_probs = 0
        # log_action_probs = F.log_softmax(action_distr, dim=1)
        # Decide on an action based on the distribution
        m = Categorical(action_probs)
        action = m.sample()

        log_prob = m.log_prob(action).unsqueeze(1)

        # break
        # Execute action in environment.
        old_state = observation

        observation, reward, done, info = env.step(action.item())
        new_state = observation

        if k % 5000 == 0:
            # print("old_state != new_state")
            # print(old_state != new_state)
            # print("oldstate " + str(old_state) + " newstate " + str(new_state))
            print("action_dist ")
            print(action_probs)
            print("On state=" + str(old_state) + ", selected action=" + str(action.item()))
            print("new state=" + str(new_state) + ", done=" + str(done) + \
                  ". Reward: " + str(reward))

        # Perform one step of the optimization
        #         policy_loss, value_loss = optimize_model(I, \
        #                                                  old_state, \
        #                                                  log_prob, \
        #                                                  log_actions_probs, \
        #                                                  action_probs, \
        #                                                  reward, \
        #                                                  new_state, \
        #                                                  entropy_impact, \
        #                                                  done)

        #         I = I * GAMMA
        # if (not done) or (done and new_state in [5,7,11,12,15]):
        memory.push(get_state_repr(old_state), action.item(), log_prob, action_probs, log_action_probs,
                    get_state_repr(new_state), reward, entropy_impact, done)

        if len(memory) >= MEMORY_SIZE:
            policy_loss, value_loss = optimize(k)
            if len(policy_loss_avg) < PRINT_OUT_TIMES:
                policy_loss_avg.append(policy_loss)
                v_loss_avg.append(value_loss)
            else:
                policy_loss_avg[episode_step % PRINT_OUT_TIMES] = policy_loss
                v_loss_avg[episode_step % PRINT_OUT_TIMES] = value_loss

        times_trained = times_trained + 1

        episode_step += 1
        # env.render()

    if k % PRINT_OUT_TIMES == 0:
        print_pi_table()
        print_v_table()

    if len(score) < 100:
        score.append(reward)
    else:
        score[k % 100] = reward

    if k % TARGET_UPDATE == 0:
        target_v_net.load_state_dict(v_net.state_dict())

    if k % PRINT_OUT_TIMES == 0:
        print("Episode {} finished after {} . Running score: {}. Policy_loss: {}, Value_loss: {}. Times trained: \
              {}. Times reached goal: {}. \
              Steps done: {}.".format(k, episode_step, np.mean(score), np.mean(policy_loss_avg), np.mean(v_loss_avg), times_trained,times_reach_goal, steps_done))
        # print("policy_loss_avg")
        # print(policy_loss_avg)
        # print("value_loss_avg")
        # print(v_loss_avg)
        #         print("times_reach_goal")
        #         print(times_reach_goal)
        times_trained = 0
        times_reach_goal = 0
        # print("Game finished. " + "-" * 5)
        # print(len(episode_series))
    #         for param in net.parameters():
    #             print(param.data)

    if reward > 0.0:
        times_reach_goal = times_reach_goal + 1


from django.urls import path

from apps.cafes.urls import CAFE_URL_KEYWORD
from apps.products import views

CATEGORY_LIST_URL_NAME = "category-list"
CATEGORY_DETAIL_URL_NAME = "category-detail"
CATEGORY_URL_KEYWORD = "category_id"

OPTION_GROUP_LIST_URL_NAME = "optiongroup-list"
OPTION_GROUP_DETAIL_URL_NAME = "optiongroup-detail"
OPTION_GROUP_URL_KEYWORD = "optiongroup_id"

PRODUCT_LIST_URL_NAME = "product-list"
PRODUCT_DETAIL_URL_NAME = "product-detail"
PRODUCT_URL_KEYWORD = "product_id"

urlpatterns = [
    path(
        f"<uuid:{CAFE_URL_KEYWORD}>/categories/",
        views.CategoryAPIViewSet.as_view({"get": "list", "post": "create"}),
        name=CATEGORY_LIST_URL_NAME,
    ),
    path(
        f"<uuid:{CAFE_URL_KEYWORD}>/categories/<int:{CATEGORY_URL_KEYWORD}>/",
        views.CategoryAPIViewSet.as_view(
            {"get": "retrieve", "put": "update", "delete": "destroy"}
        ),
        name=CATEGORY_DETAIL_URL_NAME,
    ),
    path(
        f"<uuid:{CAFE_URL_KEYWORD}>/optiongroups/",
        views.OptionGroupAPIViewSet.as_view({"get": "list", "post": "create"}),
        name=OPTION_GROUP_LIST_URL_NAME,
    ),
    path(
        f"<uuid:{CAFE_URL_KEYWORD}>/optiongroups/<int:{OPTION_GROUP_URL_KEYWORD}>/",
        views.OptionGroupAPIViewSet.as_view(
            {"get": "retrieve", "put": "update", "delete": "destroy"}
        ),
    ),
    path(
        f"<uuid:{CAFE_URL_KEYWORD}>/products/",
        views.ProductAPIViewSet.as_view({"get": "list", "post": "create"}),
        name=PRODUCT_LIST_URL_NAME,
    ),
    path(
        f"<uuid:{CAFE_URL_KEYWORD}>/products/<int:{PRODUCT_URL_KEYWORD}>/",
        views.ProductAPIViewSet.as_view(
            {"get": "retrieve", "put": "update", "delete": "destroy"}
        ),
        name=PRODUCT_DETAIL_URL_NAME,
    ),
]