import os
from collections import namedtuple
from dataclasses import dataclass, field
from datetime import datetime
from math import nan, isnan
from tempfile import NamedTemporaryFile
from typing import Union

import boto3
import pytest
from moto.s3 import mock_s3

Schema = pytest.importorskip("glue.jobs.forecast_etl").Schema
Forecast = pytest.importorskip("glue.jobs.forecast_etl").ForecastStatus
ETL = pytest.importorskip("glue.jobs.forecast_etl").ETL
ForecastDataTransformation = pytest.importorskip(
    "glue.jobs.forecast_etl"
).ForecastDataTransformation

from pyspark import SparkContext
from pyspark.sql import DataFrame
from pyspark.sql.types import StructType, StringType, StructField

from awsglue import DynamicFrame
from awsglue.context import GlueContext

DOMAINS = [
    "RETAIL",
    "CUSTOM",
    "INVENTORY_PLANNING",
    "EC2_CAPACITY",
    "WORK_FORCE",
    "METRICS",
]
TARGET_FIELDS = {
    "RETAIL": "demand",
    "CUSTOM": "target_value",
    "INVENTORY_PLANNING": "demand",
    "EC2_CAPACITY": "number_of_instances",
    "WORK_FORCE": "workforce_demand",
    "WEB_TRAFFIC": "value",
    "METRICS": "metric_value",
}
IDENTIFIERS = {
    "RETAIL": "item_id",
    "CUSTOM": "item_id",
    "INVENTORY_PLANNING": "item_id",
    "EC2_CAPACITY": "instance_type",
    "WORK_FORCE": "workforce_type",
    "WEB_TRAFFIC": "item_id",
    "METRICS": "metric_name",
}

SC: Union[SparkContext, None] = None
GC: Union[GlueContext, None] = None
SPARK_HOME_DIR = spark_home_path = os.path.abspath(
    os.path.join(
        os.path.dirname(__file__),
        os.pardir,
        os.pardir,
        os.pardir,
        ".glue",
        "spark-2.4.3-bin-spark-2.4.3-bin-hadoop2.8",
    )
)
SPARK_CONF_DIR = os.path.abspath(
    os.path.join(
        os.path.dirname(__file__),
        os.pardir,
        os.pardir,
        os.pardir,
        ".glue",
        "aws-glue-libs-glue-1.0",
        "conf",
    )
)

def glue_not_available():
    global SC
    global GC

    if SC and GC:
        return False

    if not os.path.exists(SPARK_HOME_DIR) or not os.path.exists(SPARK_CONF_DIR):
        return True

    # configuration found - try to set up Glue
    os.environ["SPARK_LOCAL_IP"] = "127.0.0.1"
    os.environ["SPARK_CONF_DIR"] = SPARK_CONF_DIR
    os.environ["SPARK_HOME"] = SPARK_HOME_DIR
    os.environ["PYSPARK_PYTHON"] = "python"
    os.environ["PYSPARK_DRIVER_PYTHON"] = "python"

    SC = SparkContext.getOrCreate()
    GC = GlueContext(SC)
    return False

def validate_data(data, expected_header, expected_values):
    if isinstance(data, ETL):
        data = data.df.toDF()
    elif isinstance(data, DynamicFrame):
        data = data.toDF()
    elif isinstance(data, DataFrame):
        pass  # already a pyspark DataFrame

    data = data.toPandas()
    columns = data.columns.to_list()
    assert columns == expected_header

    expected_rows = len(expected_values)
    expected_columns = len(expected_values[0])
    rows, columns = data.shape

    assert expected_rows == rows
    assert expected_columns == columns

    for row_idx, row in enumerate(data.values):
        for col_idx, val in enumerate(row):
            val_expected = expected_values[row_idx][col_idx]
            if isinstance(val_expected, float) and isnan(val_expected):
                assert isnan(val)
            else:
                assert val_expected == val

@pytest.fixture
def tts_data_sample():
    content = "item_id,timestamp,target_value"  # header might not exist
    content += "\nitem_01,2000-01-01 01:01:01,1"
    content += "\nitem_01,2000-01-02 02:02:02,2"
    content += "\nitem_01,2000-01-03 03:03:03,3"
    content += "\n"  # bad data - empty line
    content += "\nitem_01,2000-01-04 04:04:04,4"
    content += "\nitem_01,2000-01-05 05:05:05,5"
    content += "\nitem_01,01:01:01,5"  # bad data - bad timestamp
    content += "\nitem_01,2000-01-06 06:06:06,6"
    content += "\n,,,,,,,,,,bad_data"  # bad data, too many fields
    content += "\n,"  # bad data, not enough fields
    content += "\n,,"  # bad data, enough fields, null fields
    content += "\n,2000-01-01 01:01:01,5"  # bad data, missing identifier
    content += "\nitem_01,,5"  # bad data, missing timestamp
    content += (
        "\nitem_01,2000-01-07 07:07:07,"  # not bad data - this is a legitimate null
    )
    content += "\n\n\n"  # some extra empty lines at the end
    return content

@pytest.fixture
def md_data_sample():
    content = "item_id,geo,store_id"
    content += "\nitem_01,CA,store_01"
    content += "\nitem_02,CA,store_02"
    content += "\n\n"
    content += "\n,CA,store_02"  # bad data missing identifier
    content += "\nitem_03,CA,"  # bad data missing a metadata field
    content += "\nitem_04,,store_03"  # bad data missing a geo
    return content

@pytest.fixture
def forecast_export_sample():
    content = "item_id,date,p10,p50,p90"
    content += "\nitem_01,2020-07-01T00:00:00Z,1.0,2.0,3.0"
    return content

@pytest.fixture
def predictor_backtest_sample():
    content = "item_id,timestamp,target_value,backtestwindow_start_time,backtestwindow_end_time,p10,p50,p90"
    content += "\nitem_01,2019-07-14T12:38:26Z,0.0,2019-07-14T12:38:26,2020-06-14T12:38:26,1.0,2.0,3.0"
    return content

@pytest.fixture(params=DOMAINS)
def dataset_by_domain(request, mocker):
    dataset = mocker.MagicMock()
    fields = [
        "timestamp",
        TARGET_FIELDS[request.param],
        IDENTIFIERS[request.param],
        "geolocation",
    ]
    field_type = ["timestamp", "float", "string", "geolocation"]
    type(dataset).Schema = mocker.PropertyMock(
        return_value={
            "Attributes": [
                {"AttributeName": field, "AttributeType": field_type[idx]}
                for idx, field in enumerate(fields)
            ]
        }
    )
    return (dataset, request.param)

@pytest.fixture(params=DOMAINS)
def dynamic_frame_schema_by_domain(request):
    schema = StructType(
        [
            StructField("timestamp", StringType(), True),
            StructField("geolocation", StringType(), True),
            StructField(TARGET_FIELDS[request.param], StringType(), True),
            StructField(IDENTIFIERS[request.param], StringType(), True),
        ]
    )
    return GC.createDataFrame(SC.emptyRDD(), schema), request.param

@pytest.fixture
def forecast_service_mock(mocker):
    real_s3_cli = boto3.client("s3", region_name="us-east-1")
    cli = mocker.patch("boto3.client")
    cli().describe_dataset_group.return_value = {
        "CreationTime": datetime.now(),
        "DatasetArns": ["arn:aws:forecast:us-east-1:abcdefghijkl:dataset/test"],
        "DatasetGroupArn": "arn:aws:forecast:us-east-1:abcdefghijkl:dataset-group/test",
        "DatasetGroupName": "test",
        "Domain": "CUSTOM",
        "LastModificationTime": datetime.now(),
        "Status": "ACTIVE",
    }
    cli().describe_dataset.return_value = {
        "CreationTime": datetime.now(),
        "DataFrequency": "W",
        "DatasetArn": "arn:aws:forecast:us-east-1:abcdefghijkl:dataset/test",
        "DatasetName": "test",
        "DatasetType": "TARGET_TIME_SERIES",
        "Domain": "CUSTOM",
        "LastModificationTime": datetime.now(),
        "Schema": {
            "Attributes": [
                {"AttributeName": "item_id", "AttributeType": "string"},
                {"AttributeName": "timestamp", "AttributeType": "timestamp"},
                {"AttributeName": "target_value", "AttributeType": "float"},
            ]
        },
        "Status": "ACTIVE",
    }
    cli().describe_dataset_import_job = {
        "CreationTime": datetime.now(),
        "DatasetArn": "arn:aws:forecast:us-east-1:abcdefghijkl:dataset/test",
        "DatasetImportJobArn": "arn:aws:forecast:us-east-1:abcdefghijkl:dataset-import-job/test-2000_01_01_01_01_01",
        "DatasetImportJobName": "test-2000_01_01_01_01_01",
        "DataSize": datetime.now(),
        "DataSource": {
            "S3Config": {
                "Path": "s3://data-bucket/train.csv",
                "RoleArn": "arn:aws:iam:us-east-1:abcdefghijkl:/role/some-role",
            }
        },
        "FieldStatistics": {
            "string": {
                "Avg": 1,
                "Count": 2,
                "CountDistinct": 2,
                "CountNan": 0,
                "CountNull": 0,
                "Max": "2",
                "Min": "1",
                "Stddev": 0.5,
            }
        },
        "LastModificationTime": datetime.now,
        "Status": "ACTIVE",
    }
    cli().describe_predictor.return_value = {
        "AlgorithmArn": "arn:aws:forecast:::algorithm/ARIMA",
        "CreationTime": datetime.now(),
        "DatasetImportJobArns": [
            "arn:aws:forecast:us-east-1:abcdefghijkl:dataset-import-job/test-2000_01_01_01_01_01"
        ],
        "EvaluationParameters": {
            "BackTestWindowOffset": 3,
            "NumberOfBacktestWindows": 3,
        },
        "FeaturizationConfig": {"ForecastFrequency": "W"},
        "ForecastHorizon": 30,
        "ForecastTypes": ["0.10", "mean", "0.90"],
        "InputDataConfig": {
            "DatasetGroupArn": "arn:aws:forecast:us-east-1:abcdefghijkl:dataset-group/test",
            "SupplementaryFeatures": [{"Name": "holiday", "Value": "CA"}],
        },
        "LastModificationTime": datetime.now(),
        "PredictorArn": "arn:aws:forecast:us-east-1:abcdefghijkl:predictor/test_2000_01_01_01_01_01",
        "PredictorName": "test_2000_01_01_01_01_01",
        "Status": "ACTIVE",
    }
    cli().get_paginator.return_value.paginate.return_value = [
        {
            "Forecasts": [
                {
                    "CreationTime": datetime.now(),
                    "DatasetGroupArn": "arn:aws:forecast:us-east-1:abcdefghijkl:dataset-group/test",
                    "ForecastArn": "arn:aws:forecast:us-east-1:abcdefghijkl:forecast/test_2000_01_01_01_01_01",
                    "ForecastName": "test_2000_01_01_01_01_01",
                    "PredictorArn": "arn:aws:forecast:us-east-1:abcdefghijkl:predictor/test_2000_01_01_01_01_01",
                    "Status": "ACTIVE",
                }
            ],
            "Predictors": [
                {
                    "CreationTime": datetime.now(),
                    "DatasetGroupArn": "arn:aws:forecast:us-east-1:abcdefghijkl:dataset-group/test",
                    "LastModificationTime": datetime.now(),
                    "PredictorArn": "arn:aws:forecast:us-east-1:abcdefghijkl:predictor/test_2000_01_01_01_01_01",
                    "PredictorName": "test_2000_01_01_01_01_01",
                    "Status": "ACTIVE",
                }
            ],
            "PredictorBacktestExportJobs": [
                {
                    "CreationTime": datetime.now(),
                    "Destination": {
                        "S3Config": {
                            "KMSKeyArn": "string",
                            "Path": "s3://predictor-backtest-export/job",
                            "RoleArn": "arn:aws:iam:us-east-1:abcdefghijkl:/role/some-role",
                        }
                    },
                    "LastModificationTime": datetime.now(),
                    "PredictorBacktestExportJobArn": "arn:aws:forecast:us-east-1:abcdefghijkl:predictor-backtest-export-job/test_2000_01_01_01_01_01",
                    "PredictorBacktestExportJobName": "test_2000_01_01_01_01_01",
                    "Status": "ACTIVE",
                }
            ],
            "ForecastExportJobs": [
                {
                    "CreationTime": datetime.now(),
                    "Destination": {
                        "S3Config": {
                            "Path": "s3://forecast-export/job",
                            "RoleArn": "arn:aws:iam:us-east-1:abcdefghijkl:/role/some-role",
                        }
                    },
                    "ForecastExportJobArn": "arn:aws:forecast:us-east-1:abcdefghijkl:forecast-export-job/test_2000_01_01_01_01_01",
                    "ForecastExportJobName": "test_2000_01_01_01_01_01",
                    "LastModificationTime": datetime.now(),
                    "Message": "string",
                    "Status": "string",
                }
            ],
            "DatasetImportJobs": [
                {
                    "CreationTime": datetime.now(),
                    "DatasetImportJobArn": "arn:aws:forecast:us-east-1:abcdefghijkl:dataset-import-job/test-2000_01_01_01_01_01",
                    "DatasetImportJobName": "test-2000_01_01_01_01_01",
                    "DataSource": {
                        "S3Config": {
                            "Path": "s3://data-bucket/train.csv",
                            "RoleArn": "arn:aws:iam:us-east-1:abcdefghijkl:/role/some-role",
                        }
                    },
                    "LastModificationTime": datetime.now(),
                    "Status": "ACTIVE",
                }
            ],
        }
    ]
    return cli

def test_schema_fields(dataset_by_domain):
    dataset, domain = dataset_by_domain
    schema = Schema(dataset, domain)

    assert "geolocation" in schema.fields
    assert "timestamp" in schema.fields
    assert IDENTIFIERS[domain] in schema.fields
    assert TARGET_FIELDS[domain] in schema.fields

@pytest.mark.skipif(glue_not_available(), reason="requires local glue environment")
def test_schema_mappings(dataset_by_domain):
    dataset, domain = dataset_by_domain
    schema = Schema(dataset, domain)

    df_schema = StructType(
        [
            StructField("timestamp", StringType(), True),
            StructField(TARGET_FIELDS[domain], StringType(), True),
            StructField(IDENTIFIERS[domain], StringType(), True),
            StructField("geolocation", StringType(), True),
        ]
    )

    df = DynamicFrame.fromDF(
        dataframe=GC.createDataFrame(SC.emptyRDD(), df_schema), glue_ctx=GC, name=domain
    )
    target_field = TARGET_FIELDS[domain]
    identifier = IDENTIFIERS[domain]

    assert ("`timestamp`", "string", "timestamp", "timestamp") in schema.mappings(df)
    assert (f"`{target_field}`", "string", target_field, "double") in schema.mappings(
        df
    )
    assert (f"`{identifier}`", "string", identifier, "string") in schema.mappings(df)
    assert ("`geolocation`", "string", "geolocation", "string") in schema.mappings(df)

def test_schema_attribute_repr():
    attr = Schema._Attribute("a", "b")
    assert repr(attr) == "_Attribute('a', 'b')"

def test_forecast_domain(forecast_service_mock):
    fc = Forecast("test", "us-east-1", "abcdefghijkl")
    assert fc.domain == "CUSTOM"

def test_forecast_target_field(forecast_service_mock):
    fc = Forecast("test", "us-east-1", "abcdefghijkl")
    assert fc.target_field == "target_value"

def test_forecast_identifier_field(forecast_service_mock):
    fc = Forecast("test", "us-east-1", "abcdefghijkl")
    assert fc.identifier == "item_id"

def test_forecast_dataset_group(forecast_service_mock):
    fc = Forecast("test", "us-east-1", "abcdefghijkl")
    assert (
        fc.dataset_group.DatasetGroupArn
        == "arn:aws:forecast:us-east-1:abcdefghijkl:dataset-group/test"
    )

def test_forecast_forecast(forecast_service_mock):
    fc = Forecast("test", "us-east-1", "abcdefghijkl")
    forecast = fc.forecast
    assert (
        forecast.ForecastArn
        == "arn:aws:forecast:us-east-1:abcdefghijkl:forecast/test_2000_01_01_01_01_01"
    )

def test_forecast_predictor(forecast_service_mock):
    fc = Forecast("test", "us-east-1", "abcdefghijkl")
    predictor = fc.predictor
    assert (
        predictor.PredictorArn
        == "arn:aws:forecast:us-east-1:abcdefghijkl:predictor/test_2000_01_01_01_01_01"
    )

def test_s3_url(forecast_service_mock):
    fc = Forecast("test", "us-east-1", "abcdefghijkl")
    predictor_backtest_export_job = fc.predictor_backtest_export_job
    url = fc.s3_url(predictor_backtest_export_job)
    assert url == "s3://predictor-backtest-export/job"

    forecast_export_job = fc.forecast_export_job
    url = fc.s3_url(forecast_export_job)
    assert url == "s3://forecast-export/job"

    tts = fc._dataset_import_job("arn:aws:forecast:us-east-1:abcdefghijkl:dataset/test")
    url = fc.s3_url(tts)
    assert url == "s3://data-bucket/train.csv"

def test_target_time_series(forecast_service_mock):
    fc = Forecast("test", "us-east-1", "abcdefghijkl")
    tts = fc.target_time_series
    assert tts.DatasetArn == "arn:aws:forecast:us-east-1:abcdefghijkl:dataset/test"

def test_target_time_series_import_job(forecast_service_mock):
    fc = Forecast("test", "us-east-1", "abcdefghijkl")
    import_job = fc.target_time_series_import_job
    assert (
        import_job.DatasetImportJobArn
        == "arn:aws:forecast:us-east-1:abcdefghijkl:dataset-import-job/test-2000_01_01_01_01_01"
    )

def test_target_time_series_schema(forecast_service_mock):
    fc = Forecast("test", "us-east-1", "abcdefghijkl")
    schema = fc.target_time_series_schema
    assert schema.domain == "CUSTOM"
    assert schema.fields == ["item_id", "timestamp", "target_value"]

@mock_s3
@pytest.mark.skipif(glue_not_available(), reason="requires local glue environment")
def test_target_time_series_data(mocker, tts_data_sample, md_data_sample):
    bucket = "test-bucket"
    key_tts = "data.csv"
    key_md = "data.metadata.csv"

    cli = boto3.client("s3", region_name="us-east-1")
    cli.create_bucket(Bucket="test-bucket")
    cli.put_object(Bucket=bucket, Key=key_tts, Body=tts_data_sample)
    cli.put_object(Bucket=bucket, Key=key_md, Body=md_data_sample)

    fc = Forecast("test", "us-east-1", "abcdefghijkl")
    type(fc).predictor = mocker.PropertyMock(
        return_value=namedtuple("Predictor", "FeaturizationConfig")(
            FeaturizationConfig={"ForecastDimensions": [], "ForecastFrequency": "W"}
        )
    )
    type(fc).domain = mocker.PropertyMock(return_value="CUSTOM")
    type(fc).target_time_series_import_job = mocker.PropertyMock(
        return_value=namedtuple("DatasetImportJob", "DataSource")(
            DataSource={"S3Config": {"Path": f"s3://{bucket}/{key_tts}"}}
        )
    )
    type(fc).item_metadata_import_job = mocker.PropertyMock(
        return_value=namedtuple("DatasetImportJob", "DataSource")(
            DataSource={"S3Config": {"Path": f"s3://{bucket}/{key_md}"}}
        )
    )
    type(fc).target_time_series = mocker.PropertyMock(
        return_value=namedtuple("Dataset", "Schema")(
            Schema={
                "Attributes": [
                    {"AttributeName": "item_id", "AttributeType": "string"},
                    {"AttributeName": "timestamp", "AttributeType": "timestamp"},
                    {"AttributeName": "target_value", "AttributeType": "integer"},
                ]
            }
        )
    )
    type(fc).item_metadata = mocker.PropertyMock(
        return_value=namedtuple("Dataset", "Schema")(
            Schema={
                "Attributes": [
                    {"AttributeName": "item_id", "AttributeType": "string"},
                    {"AttributeName": "geo", "AttributeType": "string"},
                    {"AttributeName": "store_id", "AttributeType": "string"},
                ]
            }
        )
    )

    # use named temporary files to load from disk
    with NamedTemporaryFile("w", suffix="", prefix="TARGET_TIME_SERIES") as tts_data:
        with NamedTemporaryFile("w", suffix="", prefix="ITEM_METADATA") as md_data:
            tts_data.write(tts_data_sample)
            tts_data.flush()

            md_data.write(md_data_sample)
            md_data.flush()

            # point the loader to the sample data
            mocker.patch("glue.jobs.forecast_etl.TARGET_TIME_SERIES", new=tts_data.name)
            mocker.patch("glue.jobs.forecast_etl.ITEM_METADATA", new=md_data.name)
            tts = fc.target_time_series_data
            md = fc.item_metadata_data

            tts_columns_expected = ["identifier", "timestamp", "metric"]
            tts_expected = [
                ["item_01", datetime(1999, 12, 27), 3],
                ["item_01", datetime(2000, 1, 3), 18],
            ]
            validate_data(tts, tts_columns_expected, tts_expected)

            md_columns_expected = ["identifier", "geo", "store_id"]
            md_expected = [["item_01", "CA", "store_01"], ["item_02", "CA", "store_02"]]
            validate_data(md, md_columns_expected, md_expected)

def test_related_time_series(forecast_service_mock):
    fc = Forecast("test", "us-east-1", "abcdefghijkl")
    rts = fc.related_time_series
    assert not rts

def test_related_time_series_import_job(forecast_service_mock):
    fc = Forecast("test", "us-east-1", "abcdefghijkl")
    with pytest.raises(AttributeError):
        import_job = fc.related_time_series_import_job

def test_related_time_series_schema(forecast_service_mock):
    fc = Forecast("test", "us-east-1", "abcdefghijkl")
    schema = fc.related_time_series_schema
    assert not schema

def test_related_time_series_data(forecast_service_mock):
    fc = Forecast("test", "us-east-1", "abcdefghijkl")
    with pytest.raises(AttributeError):
        data = fc.related_time_series_data

def test_item_metadata(forecast_service_mock):
    fc = Forecast("test", "us-east-1", "abcdefghijkl")
    md = fc.item_metadata
    assert not md

def test_item_metadata_import_job(forecast_service_mock):
    fc = Forecast("test", "us-east-1", "abcdefghijkl")
    with pytest.raises(AttributeError):
        import_job = fc.item_metadata_import_job

def test_item_metadata_schema(forecast_service_mock):
    fc = Forecast("test", "us-east-1", "abcdefghijkl")
    schema = fc.item_metadata_schema
    assert not schema

def test_item_metadata_series_data(forecast_service_mock):
    fc = Forecast("test", "us-east-1", "abcdefghijkl")
    with pytest.raises(AttributeError):
        data = fc.item_metadata_data

@pytest.mark.skipif(glue_not_available(), reason="requires local glue environment")
def test_empty_output_schema():
    empty = Forecast.empty()
    fields = empty.schema.names

    required = [
        "identifier",
        "timestamp",
        "metric",
        *[f"p{n}" for n in range(1, 100)],
        "mean",
        *[f"dimension_{n}" for n in range(1, 11)],
        *[f"metadata_{n}" for n in range(1, 11)],
    ]
    for item in required:
        fields.remove(item)
    assert len(fields) == 0

@pytest.mark.skipif(glue_not_available(), reason="requires local glue environment")
@pytest.mark.parametrize(
    "df_a,df_b,expected_values,expected_columns",
    [
        (
            SC.parallelize([[1, 2, 3], [4, 5, 6]]).toDF(("a", "b", "c")),
            SC.parallelize([[7, 8, 9], [10, 11, 12]]).toDF(("c", "d", "a")),
            [[1, 2, 3, nan], [4, 5, 6, nan], [9, nan, 7, 8], [12, nan, 10, 11]],
            ["a", "b", "c", "d"],
        ),
        (
            SC.parallelize([[7, 8, 9], [10, 11, 12]]).toDF(("c", "d", "a")),
            SC.parallelize([[1, 2, 3], [4, 5, 6]]).toDF(("a", "b", "c")),
            [[7, 8, 9, nan], [10, 11, 12, nan], [3, nan, 1, 2], [6, nan, 4, 5]],
            ["c", "d", "a", "b"],
        ),
    ],
)
def test_union_code(df_a, df_b, expected_values, expected_columns):
    union = Forecast.union_dfs(df_a, df_b)
    validate_data(union, expected_columns, expected_values)

@mock_s3
@pytest.mark.skipif(glue_not_available(), reason="requires local glue environment")
def test_target_time_series_data(
    mocker,
    tts_data_sample,
    md_data_sample,
    predictor_backtest_sample,
    forecast_export_sample,
):
    bucket = "test-bucket"
    key_tts = "data.csv"
    key_md = "data.metadata.csv"
    key_predictor_backtest = "predictor"
    key_forecast_export = "forecast"

    cli = boto3.client("s3", region_name="us-east-1")
    cli.create_bucket(Bucket="test-bucket")
    cli.put_object(Bucket=bucket, Key=key_tts, Body=tts_data_sample)
    cli.put_object(Bucket=bucket, Key=key_md, Body=md_data_sample)
    cli.put_object(
        Bucket=bucket,
        Key=f"{key_predictor_backtest}/forecasted-values/sample.csv",
        Body=predictor_backtest_sample,
    )
    cli.put_object(
        Bucket=bucket,
        Key=f"{key_forecast_export}/sample.csv",
        Body=forecast_export_sample,
    )

    fc = Forecast("test", "us-east-1", "abcdefghijkl")
    type(fc).predictor = mocker.PropertyMock(
        return_value=namedtuple("Predictor", "FeaturizationConfig")(
            FeaturizationConfig={"ForecastDimensions": [], "ForecastFrequency": "W"}
        )
    )
    type(fc).domain = mocker.PropertyMock(return_value="CUSTOM")
    type(fc).target_time_series_import_job = mocker.PropertyMock(
        return_value=namedtuple("DatasetImportJob", "DataSource")(
            DataSource={"S3Config": {"Path": f"s3://{bucket}/{key_tts}"}}
        )
    )
    type(fc).item_metadata_import_job = mocker.PropertyMock(
        return_value=namedtuple("DatasetImportJob", "DataSource")(
            DataSource={"S3Config": {"Path": f"s3://{bucket}/{key_md}"}}
        )
    )
    type(fc).predictor_backtest_export_job = mocker.PropertyMock(
        return_value=namedtuple("PredictorBacktestExportJob", "Destination")(
            Destination={
                "S3Config": {"Path": f"s3://{bucket}/{key_predictor_backtest}"}
            }
        )
    )
    type(fc).forecast_export_job = mocker.PropertyMock(
        return_value=namedtuple("ForecastExportJob", "Destination")(
            Destination={"S3Config": {"Path": f"s3://{bucket}/{key_forecast_export}"}}
        )
    )
    type(fc).target_time_series = mocker.PropertyMock(
        return_value=namedtuple("Dataset", "Schema")(
            Schema={
                "Attributes": [
                    {"AttributeName": "item_id", "AttributeType": "string"},
                    {"AttributeName": "timestamp", "AttributeType": "timestamp"},
                    {"AttributeName": "target_value", "AttributeType": "integer"},
                ]
            }
        )
    )
    type(fc).item_metadata = mocker.PropertyMock(
        return_value=namedtuple("Dataset", "Schema")(
            Schema={
                "Attributes": [
                    {"AttributeName": "item_id", "AttributeType": "string"},
                    {"AttributeName": "geo", "AttributeType": "string"},
                    {"AttributeName": "store_id", "AttributeType": "string"},
                ]
            }
        )
    )

    # use named temporary files to load from disk
    with NamedTemporaryFile("w", suffix="", prefix="TARGET_TIME_SERIES") as tts_data:
        with NamedTemporaryFile("w", suffix="", prefix="ITEM_METADATA") as md_data:
            with NamedTemporaryFile(
                "w", suffix="", prefix="PREDICTOR_BACKTEST_EXPORT_JOB"
            ) as predictor_data:
                with NamedTemporaryFile(
                    "w", suffix="", prefix="FORECAST_EXPORT_JOB"
                ) as forecast_data:
                    tts_data.write(tts_data_sample)
                    tts_data.flush()

                    md_data.write(md_data_sample)
                    md_data.flush()

                    predictor_data.write(predictor_backtest_sample)
                    predictor_data.flush()

                    forecast_data.write(forecast_export_sample)
                    forecast_data.flush()

                    # point the loader to the sample data
                    mocker.patch(
                        "glue.jobs.forecast_etl.TARGET_TIME_SERIES", new=tts_data.name
                    )
                    mocker.patch(
                        "glue.jobs.forecast_etl.ITEM_METADATA", new=md_data.name
                    )
                    mocker.patch(
                        "glue.jobs.forecast_etl.PREDICTOR_BACKTEST_EXPORT_JOB",
                        new=predictor_data.name,
                    )
                    mocker.patch(
                        "glue.jobs.forecast_etl.FORECAST_EXPORT_JOB",
                        new=forecast_data.name,
                    )

                    # have spark load the time series data
                    tts = fc.target_time_series_data
                    md = fc.item_metadata_data
                    forecast = fc.forecast_export_job_data
                    predictor = fc.predictor_backtest_export_job_data

                    # test the individual time series are processed as expected
                    tts_columns_expected = ["identifier", "timestamp", "metric"]
                    tts_expected = [
                        ["item_01", datetime(1999, 12, 27), 3],
                        ["item_01", datetime(2000, 1, 3), 18],
                    ]
                    validate_data(tts, tts_columns_expected, tts_expected)

                    md_columns_expected = ["identifier", "geo", "store_id"]
                    md_expected = [
                        ["item_01", "CA", "store_01"],
                        ["item_02", "CA", "store_02"],
                    ]
                    validate_data(md, md_columns_expected, md_expected)

                    forecast_columns_expected = [
                        "identifier",
                        "timestamp",
                        "p10",
                        "p50",
                        "p90",
                    ]
                    forecast_expected = [
                        ["item_01", datetime(2020, 7, 1), 1.0, 2.0, 3.0]
                    ]
                    validate_data(
                        forecast, forecast_columns_expected, forecast_expected
                    )

                    predictor_columns_expected = [
                        "identifier",
                        "timestamp",
                        "metric",
                        "p10",
                        "p50",
                        "p90",
                    ]
                    predictor_expected = [
                        [
                            "item_01",
                            datetime(2019, 7, 14, 12, 38, 26),
                            0.0,
                            1.0,
                            2.0,
                            3.0,
                        ]
                    ]
                    validate_data(
                        predictor, predictor_columns_expected, predictor_expected
                    )

                    # finally aggregate the data
                    result = fc.aggregate_forecast_data.toDF().toPandas()
                    assert all(
                        identifier == "item_01"
                        for identifier in result["identifier"].to_list()
                    )
                    assert all(geo == "CA" for geo in result["metadata_1"].to_list())
                    assert all(
                        store_id == "store_01"
                        for store_id in result["metadata_2"].to_list()
                    )
                    assert result["metric"].to_list()[:3] == [3, 18, 0]
                    assert len(result) == 4

@pytest.fixture
def etl_mock():
    @dataclass
    class MockETL:
        df: Union[DataFrame, DynamicFrame, None] = field(default=None, repr=False)
        name: str = field(default="SomeName")
        gc: GlueContext = field(default=GC)
        sc: SparkContext = field(default=SC)

        def set_df(self, dataframe):
            self.df = DynamicFrame.fromDF(dataframe, GC, self.name)

    return MockETL()

@pytest.mark.skipif(glue_not_available(), reason="requires local glue environment")
def test_forecast_data_transformation_drop_duplicates(etl_mock):
    etl_mock.set_df(SC.parallelize([[1, 2, 3], [1, 2, 3]]).toDF(("a", "b", "c")))
    fdt = ForecastDataTransformation(etl_mock)
    fdt.drop_duplicates()

    validate_data(fdt.etl.df, ["a", "b", "c"], [[1, 2, 3]])


import torch
import torch.nn as nn
import torch.optim as optim
import os
import argparse
import random
import numpy as np
import time
from tqdm import tqdm

from utils import *
from sampler import *
from model import *

# comfortable path set for developing
path = '/home/doolee13/recsysData/movieLens/ml-1m/ratings.dat'
parser = argparse.ArgumentParser()
parser.add_argument('--directory', default= path)
parser.add_argument('--batch-size', default=64, type=int)
parser.add_argument('--learning-rate', default=0.001, type=float)
parser.add_argument('--maxlen', default=200, type=int)
parser.add_argument('--num-blocks', default=200, type=int)
parser.add_argument('--seed', default=42)
parser.add_argument('--hidden-units', default=50, type=int)
parser.add_argument('--dropout-rate', default=0.2, type=float)
parser.add_argument('--num-heads', default=1, type=int)
parser.add_argument('--num-epochs', default=201, type=int)
parser.add_argument('--l2-emb', default=0.0, type=float)

args = parser.parse_args()

if __name__ == '__main__':
    dataset = data_partition(args.directory)
    [train, valid, test, user_total, item_total] = dataset
    num_batch = len(train) // args.batch_size

    # seed
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.backends.cudnn.deterministic = True

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    sampler = WarpSampler(train, user_total, item_total, batch_size=args.batch_size, maxlen=args.maxlen, n_workers=3)
    # weight init done internally
    model = SASRec(user_total, item_total, device, args).to(device)
    model.train()

    criterion = nn.BCEWithLogitsLoss()
    optimizer = optim.Adam(model.parameters(), lr=args.learning_rate, betas=(0.9, 0.98))

    T = 0.0
    t0 = time.time()

    for epoch in range(1, args.num_epochs + 1):
        for step in tqdm(range(num_batch)):
            user, seq, pos, neg = sampler.next_batch()
            user, seq, pos, neg = np.array(user), np.array(seq), np.array(pos), np.array(neg)
            pos_logits, neg_logits = model(user, seq, pos, neg)
            pos_y= torch.ones(pos_logits.shape, device=device)
            neg_y= torch.zeros(neg_logits.shape, device=device)

            optimizer.zero_grad()
            indices = np.where(pos !=0)
            loss = criterion(pos_logits[indices], pos_y[indices])
            loss += criterion(neg_logits[indices], neg_y[indices])
            # l2 regularization for item embedding table necessary?
            for param in model.item_emb.parameters():
                loss += args.l2_emb * torch.norm(param)
            loss.backward()
            optimizer.step()
            print(f'loss in epoch {epoch} iteration {step} : {loss.item()}')

        if epoch % 20 == 0:
            model.eval()
            t1 = time.time() - t0
            T += t1
            print('Evaluating', end='')
            t_test = evaluate(model, dataset, args)
            t_valid = evaluate_valid(model, dataset, args)
            print('epoch:%d, time: %f(s), valid (NDCG@10: %.4f, HR@10: %.4f), test (NDCG@10: %.4f, HR@10: %.4f)'
                    % (epoch, T, t_valid[0], t_valid[1], t_test[0], t_test[1]))

            t0 = time.time()
            model.train()

    sampler.close()
    print('Done!')


# C:\Program Files\GIMP 2\lib\gimp\2.0\plug-ins

# -*- coding: utf-8 -*-

import pdb
from gimpfu import *

def draw_line(image, layer, x1, y1, x2, y2, brush_size):
    image.undo_group_start()
    pdb.gimp_context_set_foreground((255, 165, 0))
    points = [(x1, y1), (x2, y2)]
    pdb.gimp_context_set_brush_size(brush_size) # установка размера кисти
    pdb.gimp_pencil(layer, len(points), [p for point in points for p in point])
    image.undo_group_end()
    gimp.displays_flush()

# зарегистрируйте функцию в GIMP
register(
    "python_fu_draw_line",
    "Draws an orange line on the active layer",
    "Draws an orange line on the active layer",
    "Author Name",
    "Author Name",
    "2023",
    "<Image>/Filters/Draw Line...",
    "*",
    [
        (PF_INT, "x1", "X-coordinate of start point", 10),
        (PF_INT, "y1", "Y-coordinate of start point", 10),
        (PF_INT, "x2", "X-coordinate of end point", 1250),
        (PF_INT, "y2", "Y-coordinate of end point", 1250),
        (PF_INT, "brush_size", "Brush size", 500)
    ],
    [],
    draw_line)

# запустите GIMP
main()


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu May 11 11:07:56 2023

@author: lucas
"""

import streamlit as st
from datetime import date, timedelta
import re
import yfinance as yf
from plotly import graph_objs as go
from prophet.forecaster import Prophet
from prophet.plot import plot_plotly

st.title("ðŸ’°ðŸš€ Stock Predictor App ðŸš€ðŸ’°")

# Set the minimum and maximum start date values
min_date = date(2010, 1, 1)
max_date = date.today() - timedelta(days=1)

# Get the start date from the user using a slider
START= st.slider(
    'Select a start date:',
    min_value=min_date,
    max_value=max_date,
    value=date(2015, 1, 1),
    format="YYYY-MM-DD"
)

# Print the selected start date
st.write('Start date selected:', START.strftime('%Y-%m-%d'))

TODAY = date.today().strftime('%Y-%m-%d')

# Function to assess if the Input for a Ticker is valid
def is_valid_ticker(ticker):
    """Check if the ticker is valid."""
    pattern = r'^[A-Z.]{1,6}$'  # match 1 to 5 uppercase letters
    return re.match(pattern, ticker)

# Check if user_input is valid and return a warning message otherwise
user_input = st.text_input('Enter Stock Ticker', 'AAPL')
if not is_valid_ticker(user_input):
    st.warning('Please enter a valid stock ticker (e.g. AAPL)')

n_years = st.slider('Years of prediction:', 1, 4)
period = n_years*365

@st.cache_data
def load_data(ticker, START):
    data = yf.download(ticker, start=START, end=TODAY, repair=True)
    data.reset_index(inplace=True)
    return data

data_load_state = st.text(f"Loading data for: {user_input} from {START} to {TODAY}...")
data = load_data(user_input, START)
data_load_state.text(f'Loading data for: {user_input} from {START} to {TODAY} is done!')

# Get Ticker object for the given ticker
ticker_info = yf.Ticker(user_input)

# Fetch company summary
company_summary = ticker_info.info['longBusinessSummary']

# Display the company summary
st.subheader('Company Summary:')
st.write(company_summary)

st.subheader('Raw Data')
st.write(data.head())

def plot_raw_data():
    fig = go.Figure()
    fig.add_trace(go.Scatter(x=data['Date'], y=data['Open'], name='stock_open'))
    fig.add_trace(go.Scatter(x=data['Date'], y=data['Close'], name='stock_close'))
    fig.layout.update(title_text='Time Series Data', xaxis_rangeslider_visible=True)
    st.plotly_chart(fig)

plot_raw_data()

# Forecasting
df_train = data[['Date', 'Close']]
df_train = df_train.rename(columns={'Date': 'ds', 'Close':'y'})

m = Prophet()
m.fit(df_train)
future = m.make_future_dataframe(periods=period)
forecast = m.predict(future)

st.subheader('Forecast data')
st.write(data.tail())

st.write('forecast data')
fig1 = plot_plotly(m, forecast)
st.plotly_chart(fig1)

st.write('forecast components')
fig2 = m.plot_components(forecast)
st.write(fig2)


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from keras.layers import Dense, Input
from keras.layers import Conv2D, Flatten, Lambda
from keras.layers import Reshape, Conv2DTranspose
from keras.models import Model
from keras.losses import mse, binary_crossentropy
from keras.utils import plot_model
from keras import backend as K

from sklearn.model_selection import train_test_split

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

K.clear_session()

np.random.seed(237)

# reparameterization trick
# instead of sampling from Q(z|X), sample eps = N(0,I)
# then z = z_mean + sqrt(var)*eps
def sampling(args):
    """Reparameterization trick by sampling fr an isotropic unit Gaussian.
    # Arguments
        args (tensor): mean and log of variance of Q(z|X)
    # Returns
        z (tensor): sampled latent vector
    """

    z_mean, z_log_var = args
    batch = K.shape(z_mean)[0]
    dim = K.int_shape(z_mean)[1]
    # by default, random_normal has mean=0 and std=1.0
    epsilon = K.random_normal(shape=(batch, dim))
    return z_mean + K.exp(0.5 * z_log_var) * epsilon

def process_data(data_path):

    data = np.load(data_path)

    X_train, X_test = train_test_split(data, test_size=0.05, random_state=42)
    print('Shape train/test:', X_train.shape, X_test.shape)

    image_size = X_train.shape[1], X_train.shape[2]

    data = np.reshape(data, [-1, image_size[0], image_size[1], 1])
    X_train = np.reshape(X_train, [-1, image_size[0], image_size[1], 1])
    X_test = np.reshape(X_test, [-1, image_size[0], image_size[1], 1])

    data = data.astype('float32') / 255
    X_train = X_train.astype('float32') / 255
    X_test = X_test.astype('float32') / 255

    return data, X_train, X_test, image_size

def construct_vae(image_size, kernel_size, latent_dim):
    # network parameters
    input_shape = (image_size[0], image_size[1], 1)

    # VAE model = encoder + decoder
    # build encoder model
    inputs = Input(shape=input_shape, name='encoder_input')
    x = inputs
    x = Conv2D(filters=16, kernel_size=kernel_size, activation='relu', strides=1, padding='same')(x)
    x = Conv2D(filters=32, kernel_size=kernel_size, activation='relu', strides=2, padding='same')(x)
    x = Conv2D(filters=64, kernel_size=kernel_size, activation='relu', strides=1, padding='same')(x)

    # shape info needed to build decoder model
    shape = K.int_shape(x)

    # generate latent vector Q(z|X)
    x = Flatten()(x)
    x = Dense(16, activation='relu')(x)
    z_mean = Dense(latent_dim, name='z_mean')(x)
    z_log_var = Dense(latent_dim, name='z_log_var')(x)

    # use reparameterization trick to push the sampling out as input
    # note that "output_shape" isn't necessary with the TensorFlow backend
    z = Lambda(sampling, output_shape=(latent_dim,), name='z')([z_mean, z_log_var])

    # instantiate encoder model
    encoder = Model(inputs, [z_mean, z_log_var, z], name='encoder')
    encoder.summary()
    plot_model(encoder, to_file='vae_cnn_encoder.png', show_shapes=True)

    # build decoder model
    latent_inputs = Input(shape=(latent_dim,), name='z_sampling')
    x = Dense(shape[1] * shape[2] * shape[3], activation='relu')(latent_inputs)
    x = Reshape((shape[1], shape[2], shape[3]))(x)

    x = Conv2DTranspose(filters=64, kernel_size=kernel_size, activation='relu', strides=1, padding='same')(x)
    x = Conv2DTranspose(filters=32, kernel_size=kernel_size, activation='relu', strides=2, padding='same')(x)
    x = Conv2DTranspose(filters=16, kernel_size=kernel_size, activation='relu', strides=1, padding='same')(x)

    outputs = Conv2DTranspose(filters=1,
                              kernel_size=kernel_size,
                              activation='sigmoid',
                              padding='same',
                              name='decoder_output')(x)

    # instantiate decoder model
    decoder = Model(latent_inputs, outputs, name='decoder')
    decoder.summary()
    plot_model(decoder, to_file='vae_cnn_decoder.png', show_shapes=True)

    # instantiate VAE model
    outputs = decoder(encoder(inputs)[2])
    vae = Model(inputs, outputs, name='vae')

    # VAE loss = mse_loss or xent_loss + kl_loss
    reconstruction_loss = binary_crossentropy(K.flatten(inputs), K.flatten(outputs))

    reconstruction_loss *= image_size[0] * image_size[1]
    kl_loss = 1 + z_log_var - K.square(z_mean) - K.exp(z_log_var)
    kl_loss = K.sum(kl_loss, axis=-1)
    kl_loss *= -0.5
    vae_loss = K.mean(reconstruction_loss + kl_loss)
    vae.add_loss(vae_loss)
    vae.compile(optimizer='rmsprop')
    vae.summary()
    plot_model(vae, to_file='vae_cnn.png', show_shapes=True)

    return vae, encoder, decoder

if __name__ == '__main__':

    is_train = False
    data_file = '../data/out/moment_frames_5.npy'
    data, X_train, X_test, im_size = process_data(data_file)

    kernel_size = (3, 3)
    latent_dim = 128
    batch_size = 128
    epochs = 10

    vae, encoder, decoder = construct_vae(im_size, kernel_size, latent_dim)

    if is_train:
        history = vae.fit(X_train,
                          epochs=epochs,
                          batch_size=batch_size,
                          validation_data=(X_test, None),
                          verbose=2)
        vae.save_weights('vae_cnn.h5')

        # summarize history for loss
        plt.plot(history.history['loss'])
        plt.plot(history.history['val_loss'])
        plt.title('model loss')
        plt.ylabel('loss')
        plt.xlabel('epoch')
        plt.legend(['train', 'test'], loc='upper left')
        plt.savefig('vae_train.jpeg')
        plt.show()

    else:
        vae.load_weights('vae_cnn.h5')

    # Transform to latent representation
    encoded_data = encoder.predict(data, batch_size=batch_size)

    pd.DataFrame(encoded_data[0]).to_csv('latest_rep_cnn.csv', index=None)

    print('Completed.')


# @Project : ftry
# @Filename: textocr
# @Date    : 2017-03-03
# @Author  : Shiyue Nie

# ref : https://cloud.google.com/vision/docs/how-to
#       https://cloud.google.com/vision/docs/tutorials

import io
import os

# Imports the Google Cloud client library
from google.cloud import vision

# from oauth2client.client import GoogleCredentials
# os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = "E:\DHLAB\Google_Vision\ptry-efe4ae49a335.json"
# credentials = vision.Credentials()
# credentials = GoogleCredentials.get_application_default()

def detect_label(path):
    """Detects label in the file."""

    # Loads the image into memory
    with io.open(path, 'rb') as image_file:
        content = image_file.read()

    image = vision_client.image(content=content)
    # Performs label detection on the image file

    labels = image.detect_labels()
    print('Labels:')
    for label in labels:
        print(label.description)

def detect_text(path):
    """Detects text in the file."""

    with io.open(path, 'rb') as image_file:
        content = image_file.read()

    image = vision_client.image(content=content)

    texts = image.detect_text()
    print('Texts:')
    #for text in texts:
    #    print(text.description)
    return texts[0].description

# service account is stored in a JSON keyfile.
# Instantiates a client
vision_client = vision.Client().from_service_account_json('E:\DHLAB\Google_Vision\ptry-efe4ae49a335.json')
# vision_client = vision.Client(credentials=creds)

# The name of the image file to annotate
file_name = os.path.join(
    os.path.dirname(__file__),
    'E:/workplace/pycharm/ftry/test/02.png')

# detect_label(file_name)
print(detect_text(file_name))
#detect_text(file_name)


# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this file,
# You can obtain one at http://mozilla.org/MPL/2.0/.

import datetime
import os
import shutil
import subprocess
import tarfile
import tempfile
from email.utils import format_datetime
from pathlib import Path
from string import Template

import mozpack.path as mozpath
from mozpack.files import FileFinder

from mozbuild.repackaging.application_ini import get_application_ini_values

class NoPackageFound(Exception):
    """Raised when no .deb is found after calling dpkg-buildpackage"""

_DEB_ARCH = {
    "x86": "i386",
    "x86_64": "amd64",
}
# At the moment the Firefox build baseline is jessie.
# The debian-repackage image defined in taskcluster/docker/debian-repackage/Dockerfile
# bootstraps the /srv/jessie-i386 and /srv/jessie-amd64 chroot environments we use to
# create the `.deb` repackages. By running the repackage using chroot we generate shared
# library dependencies that match the Firefox build baseline
# defined in taskcluster/scripts/misc/build-sysroot.sh
_DEB_DIST = "jessie"

def repackage_deb(infile, output, template_dir, arch):

    if not tarfile.is_tarfile(infile):
        raise Exception("Input file %s is not a valid tarfile." % infile)

    deb_arch = _DEB_ARCH[arch]

    if os.path.isdir(f"/srv/{_DEB_DIST}-{deb_arch}"):
        tmpdir = tempfile.mkdtemp(dir=f"/srv/{_DEB_DIST}-{deb_arch}/tmp")
    else:
        tmpdir = tempfile.mkdtemp()

    extract_dir = os.path.join(tmpdir, "source")
    try:
        with tarfile.open(infile) as tar:
            tar.extractall(path=extract_dir)
        finder = FileFinder(extract_dir)
        values = get_application_ini_values(
            finder,
            dict(section="App", value="Name"),
            dict(section="App", value="CodeName", fallback="Name"),
            dict(section="App", value="Vendor"),
            dict(section="App", value="RemotingName"),
            dict(section="App", value="Version"),
            dict(section="App", value="BuildID"),
        )
        app_name = next(values)
        displayname = next(values)
        vendor = next(values)
        remotingname = next(values)
        version = next(values)
        buildid = next(values)
        timestamp = datetime.datetime.strptime(buildid, "%Y%m%d%H%M%S")

        os.mkdir(mozpath.join(extract_dir, "debian"))
        shutil.copy(
            mozpath.join(template_dir, "rules"),
            mozpath.join(extract_dir, "debian", "rules"),
        )
        defines = {
            "DEB_DESCRIPTION": f"{vendor} {displayname}",
            "DEB_PKG_NAME": remotingname.lower(),
            "DEB_PKG_VERSION": f"{version}.{buildid}",
            "DEB_CHANGELOG_DATE": format_datetime(timestamp),
        }

        template_dir_filenames = os.listdir(template_dir)
        template_filenames = [
            filename for filename in template_dir_filenames if filename.endswith(".in")
        ]
        plain_filenames = [
            filename
            for filename in template_dir_filenames
            if not filename.endswith(".in")
        ]

        for filename in template_filenames:
            with open(mozpath.join(template_dir, filename)) as f:
                template = Template(f.read())
            with open(
                mozpath.join(extract_dir, "debian", Path(filename).stem), "w"
            ) as f:
                f.write(template.substitute(defines))

        for filename in plain_filenames:
            shutil.copy(
                mozpath.join(template_dir, filename),
                mozpath.join(extract_dir, "debian"),
            )

        with open(
            mozpath.join(extract_dir, app_name.lower(), "is-packaged-app"), "w"
        ) as f:
            f.write("This is a packaged app.\n")

        if os.path.isdir(f"/srv/{_DEB_DIST}-{deb_arch}"):
            subprocess.check_call(
                [
                    "chroot",
                    f"/srv/{_DEB_DIST}-{deb_arch}",
                    "bash",
                    "-c",
                    f"cd /tmp/*/source; dpkg-buildpackage -us -uc -b -a{deb_arch}",
                ]
            )
        else:
            # If there's is no _DEB_DIST bootstrapped in /srv run the packaging here.
            subprocess.check_call(
                ["dpkg-buildpackage", "-us", "-uc", "-b", f"-a{deb_arch}"],
                cwd=extract_dir,
            )

        deb_file_name = (
            f"{defines['DEB_PKG_NAME']}_{defines['DEB_PKG_VERSION']}_{deb_arch}.deb"
        )
        deb_file_path = mozpath.join(tmpdir, deb_file_name)

        if not os.path.exists(deb_file_path):
            raise NoPackageFound(
                f"No {deb_file_name} package found after calling dpkg-buildpackage"
            )

        subprocess.check_call(["dpkg-deb", "--info", deb_file_path])

        shutil.move(deb_file_path, output)

    finally:
        shutil.rmtree(tmpdir)


import numpy as np
import pandas as pd
from visual_behavior.data_access import reformat
import visual_behavior_glm.GLM_fit_tools as gft
import visual_behavior_glm.GLM_visualization_tools as gvt
import visual_behavior_glm.build_dataframes as bd
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import cross_val_score
from sklearn.model_selection import cross_val_predict
from sklearn.model_selection import cross_validate
import matplotlib.pyplot as plt
from tqdm import tqdm

def decode_experiment(oeid, data='events',window=[0,.75]):

    # Load SDK object
    print('Loading data')
    run_params = {'include_invalid_rois':False}
    session = gft.load_data(oeid, run_params)
    reformat.add_licks_each_flash(session.stimulus_presentations, session.licks)
    reformat.add_rewards_each_flash(session.stimulus_presentations, session.rewards)
    session.stimulus_presentations['licked'] = [True if len(licks) > 0 else False \
        for licks in session.stimulus_presentations.licks.values]
    session.stimulus_presentations['hit'] =  \
        session.stimulus_presentations['licked'] & \
        session.stimulus_presentations['is_change']
    session.stimulus_presentations['miss'] = \
        ~session.stimulus_presentations['licked'] & \
        session.stimulus_presentations['is_change']
    session.stimulus_presentations.at[\
        ~session.stimulus_presentations['is_change'],'hit'] = np.nan
    session.stimulus_presentations.at[\
        ~session.stimulus_presentations['is_change'],'miss'] = np.nan

    # get list of cell dataframes
    print('Generating cell dataframes')
    cells = get_cells(session,data=data,window=window)

    # Perform decoding iterating over the number of cells
    print('Decoding')
    results_df = iterate_n_cells(cells)

    # Save results
    print('Save results')
    filename='/allen/programs/braintv/workgroups/nc-ophys/alex.piet'+\
        '/behavior/decoding/experiments/'
    filename += str(oeid)+'.pkl'
    print('Saving to: '+filename)
    results_df.to_pickle(filename)
    print('Finished')

def load_experiment_results(oeid,version=None):
    if version is not None:
        filename='/allen/programs/braintv/workgroups/nc-ophys/alex.piet'+\
            '/behavior/decoding/experiments_fit_{}/'.format(version)
    else:
        filename='/allen/programs/braintv/workgroups/nc-ophys/alex.piet'+\
            '/behavior/decoding/experiments/'
    filename += str(oeid)+'.pkl'
    return pd.read_pickle(filename)

def load_all(experiment_table,summary_df,version=None,mesoscope_only=False):

    # Iterate through experiments and load decoding results
    dfs = []
    failed = 0
    if mesoscope_only:
        summary_df = summary_df.query('equipment_name == "MESO.1"').copy()
    oeids = np.concatenate(summary_df['ophys_experiment_id'].values)

    for oeid in oeids:
        try:
            df = load_experiment_results(oeid,version)
        except:
            failed +=1
        else:
            df['ophys_experiment_id'] = oeid
            dfs.append(df)
    print('Failed to load: {}'.format(failed))

    # Merge and add experiment meta data
    print('Concatenating')
    df = pd.concat(dfs,sort=True)
    print('merging experiment table')
    df = pd.merge(df,experiment_table.reset_index()[['ophys_experiment_id',\
        'ophys_session_id','behavior_session_id','targeted_structure',\
        'imaging_depth','equipment_name']],on='ophys_experiment_id')
    print('merging summary_df')
    df = pd.merge(df, summary_df[['ophys_session_id','visual_strategy_session',\
        'experience_level','cre_line']],on='ophys_session_id')

    return df

def get_cells(session, data='events',window=[0,.75]):
    '''
        Iterate over all cells in this experiment and make a list
        of cell dataframes
    '''

    # Iterate over all cells in experiment
    cells = []
    cell_specimen_ids = session.cell_specimen_table.index.values
    for cell in tqdm(cell_specimen_ids):
        # Generate the dataframe for this cell
        cell_df = get_cell_table(session, cell,data=data,window=window)
        cells.append(cell_df)

    # return list of cell dataframes
    return cells

def get_cell_table(session, cell_specimen_id, data='events',
    window=[0,.75],balance=True):
    '''
        Generates a dataframe for one cell where rows are image presentations
        and the response at each timepoint is a column
    '''
    # Get cell activity interpolated onto constant time points
    df = bd.get_cell_df(session, cell_specimen_id, data=data)
    full_df = bd.get_cell_etr(df, session, time = window)

    # Pivot the table such that all responses at the same time point are a column
    cell = pd.pivot_table(full_df, values='response',
        index='stimulus_presentations_id', columns=['time'])

    # Annotate stimulus information
    cell = pd.merge(cell, session.stimulus_presentations,
        left_index=True, right_index=True)

    # Balance classes by using just changes and the image before
    cell['pre_change'] = cell['is_change'].shift(-1,fill_value=False)
    if balance:
        cell = cell.query('is_change or pre_change')

    return cell

def get_matrix(cell):
    '''
        Grabs the responses of the cell at each time point across each
        image presentation and returns a numpy matrix
    '''
    cols = np.sort([x for x in cell.columns.values if not isinstance(x,str)])
    x = cell[cols].to_numpy()
    return x

def plot_by_strategy_performance(visual, timing,savefig,cell_type):
    # Plot decoder performance
    plt.figure()
    plt.clf()

    # visual decoder performance
    v = visual.groupby('n_cells')
    summary = v.mean()
    summary['size'] = v.size()
    summary['sem_score'] = v['test_score'].sem()
    plt.errorbar(summary.index.values, summary['test_score'],
        yerr=summary['sem_score'],color='darkorange')
    plt.plot(summary.index.values, summary.test_score,'o-',color='darkorange',
        label='visual sessions')

    #timing decoder performance
    t = timing.groupby('n_cells')
    summary = t.mean()
    summary['size'] = t.size()
    summary['sem_score'] = t['test_score'].sem()
    plt.errorbar(summary.index.values, summary['test_score'],
        yerr=summary['sem_score'],color='blue')
    plt.plot(summary.index.values, summary.test_score,'bo-',
        label='timing sessions')

    # Clean up decoder performance plot
    plt.xlim(0,80)
    plt.ylim(0.5,1)
    plt.ylabel('decoder performance',fontsize=16)
    plt.xlabel('number of cells',fontsize=16)
    ax = plt.gca()
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.xaxis.set_tick_params(labelsize=12)
    ax.yaxis.set_tick_params(labelsize=12)
    plt.legend(loc='upper right')
    plt.tight_layout()

    # save decoder performance figure
    if savefig:
        filename='/allen/programs/braintv/workgroups/nc-ophys/alex.piet'+\
            '/behavior/decoding/figures/'
        filename += cell_type+'_decoder_performance.png'
        plt.savefig(filename)

def plot_by_strategy_correlation(visual, timing, savefig, cell_type):
    # Plot behavior correlation figure
    plt.figure()
    plt.clf()

    # visual correlation
    v = visual.groupby('n_cells')
    summary = v.mean()
    summary['size'] = v.size()
    summary['sem_behavior_correlation'] = v['behavior_correlation'].sem()
    plt.errorbar(summary.index.values, summary['behavior_correlation'],
        yerr=summary['sem_behavior_correlation'],color='darkorange')
    plt.plot(summary.index.values, summary.behavior_correlation,'o-',
        color='darkorange', label='visual sessions')

    # timing correlation
    t = timing.groupby('n_cells')
    summary = t.mean()
    summary['size'] = t.size()
    summary['sem_behavior_correlation'] = t['behavior_correlation'].sem()
    plt.errorbar(summary.index.values, summary['behavior_correlation'],
        yerr=summary['sem_behavior_correlation'], color='blue')
    plt.plot(summary.index.values, summary.behavior_correlation,'bo-',
        label='timing sessions')

    # clean up correlation figure
    plt.xlim(0,80)
    plt.ylim(0,0.5)
    plt.ylabel('decoder correlation with behavior',fontsize=16)
    plt.xlabel('number of cells',fontsize=16)
    ax = plt.gca()
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.xaxis.set_tick_params(labelsize=12)
    ax.yaxis.set_tick_params(labelsize=12)
    plt.legend(loc='upper right')
    plt.tight_layout()

    # save behavior correlation figure
    if savefig:
        filename='/allen/programs/braintv/workgroups/nc-ophys/alex.piet'+\
            '/behavior/decoding/figures/'
        filename += cell_type+'_decoder_correlation.png'
        plt.savefig(filename)

def plot_by_strategy_scatter(visual, timing, metric, savefig, cell_type,
    version=None,ax=None,meso=False):

    # Plot behavior correlation figure
    if ax is None:
        fig, ax = plt.subplots()
        ax.set_aspect('equal')
    colors = gvt.project_colors()
    color = colors[cell_type]

    # visual correlation
    v = visual.groupby('n_cells')
    vsummary = v.mean()
    vsummary['size'] = v.size()
    vsummary['sem_'+metric] = v[metric].sem()

    t = timing.groupby('n_cells')
    tsummary = t.mean()
    tsummary['size'] = t.size()
    tsummary['sem_'+metric] = t[metric].sem()

    n = np.min([len(vsummary), len(tsummary)])

    ax.errorbar(tsummary.iloc[0:n][metric],
        vsummary.iloc[0:n][metric],
        xerr=tsummary.iloc[0:n]['sem_'+metric],
        yerr=vsummary.iloc[0:n]['sem_'+metric],
        color=color,fmt='o',markersize=1,alpha=.5)

    for index in range(0,n):
        ax.scatter(tsummary.iloc[index][metric],
            vsummary.iloc[index][metric],
            s=tsummary.index.values[index]*1.5,color=color,
            label=tsummary.index.values[index],alpha=1)

    if metric == 'behavior_correlation':
        ax.plot([0, .2],[0, .2], 'k--',alpha=.25)
        ax.set_ylabel('visual sessions',
            fontsize=16)
        ax.set_xlabel('timing sessions',
            fontsize=16)
        if meso:
            ax.set_xlim(0,.2)
            ax.set_ylim(0,.2)
            ticks = np.arange(0,.22,.04)
        else:
            ax.set_xlim(0,.165)
            ax.set_ylim(0,.165)
            ticks = np.arange(0,.18,.02)
        ax.set_xticks(ticks)
        ax.set_yticks(ticks)
        ax.set_title('change decoder correlation with behavior',fontsize=16)
    elif metric == 'test_score':
        ax.plot([0.5, 1],[0.5, 1], 'k--',alpha=.25)
        ax.set_ylabel('visual sessions',
            fontsize=16)
        ax.set_xlabel('timing sessions',
            fontsize=16)
        ax.set_xlim(0.5,.75)
        ax.set_ylim(0.5,.75)
        ax.set_title('change decoder performance',fontsize=16)
    elif metric == 'test_score_hit_vs_miss':
        ax.plot([0.5, 1],[0.5, 1], 'k--',alpha=.25)
        ax.set_ylabel('visual sessions',
            fontsize=16)
        ax.set_xlabel('timing sessions',
            fontsize=16)
        if meso:
            ax.set_xlim(0.5,.8)
            ax.set_ylim(0.5,.8)
        else:
            ax.set_xlim(0.5,.75)
            ax.set_ylim(0.5,.75)
        ax.set_title('hit decoder performance',fontsize=16)
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.xaxis.set_tick_params(labelsize=12)
    ax.yaxis.set_tick_params(labelsize=12)

    ax.set_aspect('equal')
    ax.legend(loc='upper left',bbox_to_anchor=(1.04,1),title='# of cells')
    plt.tight_layout()

    if savefig:
        if version is None:
            filename='/allen/programs/braintv/workgroups/nc-ophys/alex.piet'+\
                '/behavior/decoding/figures/'
        else:
            filename='/allen/programs/braintv/workgroups/nc-ophys/alex.piet'+\
                '/behavior/decoding/figures_fit_{}/'.format(version)
        if meso:
            filename += 'scatter_by_cre_meso_'+metric+'.svg'
        else:
            filename += 'scatter_by_cre_'+metric+'.svg'
        print(filename)
        plt.savefig(filename)

def plot_by_strategy_hit_vs_miss(visual, timing, savefig, cell_type):

    # Plot hit vs miss decoder performance
    plt.figure()
    plt.clf()

    # visual hit vs miss decoder performance
    v = visual.groupby('n_cells')
    summary = v.mean()
    summary['size'] = v.size()
    summary['sem_score'] = v['test_score_hit_vs_miss'].sem()
    plt.errorbar(summary.index.values, summary['test_score_hit_vs_miss'],
        yerr=summary['sem_score'],color='darkorange')
    plt.plot(summary.index.values, summary.test_score_hit_vs_miss,'o-',
        color='darkorange',label='visual sessions')

    #timing hit vs miss decoder performance
    t = timing.groupby('n_cells')
    summary = t.mean()
    summary['size'] = t.size()
    summary['sem_score'] = t['test_score_hit_vs_miss'].sem()
    plt.errorbar(summary.index.values, summary['test_score_hit_vs_miss'],
        yerr=summary['sem_score'],color='blue')
    plt.plot(summary.index.values, summary.test_score_hit_vs_miss,'bo-',
        label='timing sessions')

    # Clean up hit vs miss decoder performance plot
    plt.xlim(0,80)
    plt.ylim(0.5,1)
    plt.ylabel('hit vs miss decoder performance',fontsize=16)
    plt.xlabel('number of cells',fontsize=16)
    ax = plt.gca()
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.xaxis.set_tick_params(labelsize=12)
    ax.yaxis.set_tick_params(labelsize=12)
    plt.legend(loc='upper right')
    plt.tight_layout()

    # save hit vs miss decoder performance figure
    if savefig:
        filename='/allen/programs/braintv/workgroups/nc-ophys/alex.piet'+\
            '/behavior/decoding/figures/'
        filename += cell_type+'_hit_vs_miss_decoder_performance.png'
        plt.savefig(filename)

def plot_by_strategy(results_df,aggregate_first=True,cell_type='exc',
    areas=['VISp','VISl'],equipment=['MESO.1','CAM2P.3','CAM2P.4','CAM2P.5'],
    savefig=False):

    # parse cell type
    mapper = {
        'exc':'Slc17a7-IRES2-Cre',
        'sst':'Sst-IRES-Cre',
        'vip':'Vip-IRES-Cre'
        }
    cre = mapper[cell_type]

    # filter out experiments
    results_df = results_df.query('experience_level == "Familiar"')
    results_df = results_df.query('cre_line == @cre')
    results_df = results_df.query('targeted_structure in @areas')
    results_df = results_df.query('equipment_name in @equipment')

    # Average over samples from the same experiment, so each experiment
    # is weighted the same
    if aggregate_first:
        x = results_df.groupby(['n_cells','ophys_experiment_id']).mean()
        results_df = x.reset_index()

    # Split by strategy
    visual = results_df.query('visual_strategy_session')
    timing = results_df.query('not visual_strategy_session')

    plot_by_strategy_performance(visual, timing,savefig,cell_type)
    plot_by_strategy_correlation(visual, timing, savefig, cell_type)
    plot_by_strategy_hit_vs_miss(visual, timing, savefig, cell_type)
    plot_by_strategy_scatter(visual, timing,'behavior_correlation',
        savefig, cell_type,version)
    plot_by_strategy_scatter(visual, timing,'test_score',savefig,
        cell_type,version)

def plot_by_cre(results_df, aggregate_first=True,areas=['VISp','VISl'],
    equipment=['MESO.1','CAM2P.3','CAM2P.4','CAM2P.5'],ncells=2,
    savefig=False,version=None,meso=False):

    # filter out experiments
    results_df = results_df.query('experience_level == "Familiar"')
    results_df = results_df.query('targeted_structure in @areas')
    results_df = results_df.query('equipment_name in @equipment')
    results_df = results_df.query('n_cells > @ncells')

    # Split by strategy
    cres = ['Slc17a7-IRES2-Cre','Sst-IRES-Cre','Vip-IRES-Cre']
    mapper = {
        'Slc17a7-IRES2-Cre':'exc',
        'Sst-IRES-Cre':'sst',
        'Vip-IRES-Cre':'vip'
        }
    fig1, ax1 = plt.subplots()
    fig2, ax2 = plt.subplots()
    fig3, ax3 = plt.subplots()
    for index,cre in enumerate(cres):
        cre_df = results_df.query('cre_line == @cre')
        # Average over samples from the same experiment, so each experiment
        # is weighted the same
        if aggregate_first:
            x = cre_df.groupby(['n_cells','ophys_experiment_id']).mean()
            cre_df = x.reset_index()

        visual = cre_df.query('(visual_strategy_session)')
        timing = cre_df.query('(not visual_strategy_session)')
        if index == 2:
            plt.figure(fig1.number)
            plot_by_strategy_scatter(visual,timing,'behavior_correlation',
                savefig,mapper[cre],version,ax1,meso=meso)
            plt.figure(fig2.number)
            plot_by_strategy_scatter(visual,timing,'test_score',
                savefig,mapper[cre],version,ax2,meso=meso)
            plt.figure(fig3.number)
            plot_by_strategy_scatter(visual,timing,'test_score_hit_vs_miss',
                savefig,mapper[cre],version,ax3,meso=meso)
        else:
            plot_by_strategy_scatter(visual,timing,'behavior_correlation',
                False,mapper[cre],version,ax1,meso=meso)
            plot_by_strategy_scatter(visual,timing,'test_score',
                False,mapper[cre],version,ax2,meso=meso)
            plot_by_strategy_scatter(visual,timing,'test_score_hit_vs_miss',
                False,mapper[cre],version,ax3,meso=meso)

def iterate_n_cells(cells):
    n_cells = [1,2,5,10,20,40,80]

    results = {}
    for n in n_cells:
        if len(cells) < n:
            break
        print('Decoding with n={} cells'.format(n))
        results[n] = decode_cells(cells, n)

    results_df = pd.concat(results)
    return results_df

def decode_cells(cells, n_cells):
    '''
        Cells is a list of dataframes, one for each cell
        n_cells is the number of cells to decode with in each sample
    '''

    # How many times we need to sample in order to get 99% chance each cell
    # is used at least once
    if n_cells == 1:
        n_samples = len(cells)
    else:
        n_samples = int(np.ceil(np.log(0.01)/np.log(1-n_cells/len(cells))))

    # Iterate over samples and save output
    output = []
    for n in tqdm(range(0,n_samples)):
        if n_cells == 1:
            temp = decode_cells_sample(cells, n_cells, index=n)
        else:
            temp = decode_cells_sample(cells, n_cells)
        output.append(temp)

    # Return the output of the samples
    output_df = pd.DataFrame(output)
    output_df['n_cells'] = n_cells
    return output_df

def decode_cells_sample(cells, n_cells,index=None):
    '''
        Sample from the list of cells, and perform decoding once
        returns the output of the decoding
    '''

    # Sample n_cells from list of cells
    if n_cells == 1:
        sample_cells = [cells[index]]
    else:
        cells_in_sample = np.random.choice(len(cells),n_cells, replace=False)
        sample_cells = [cells[i] for i in cells_in_sample]

    # Construct X
    X = []
    for cell in sample_cells:
        X.append(get_matrix(cell))
    X = np.concatenate(X,axis=1)

    # y is the same for every cell, since its a behavioral output
    y = sample_cells[0]['is_change'].values

    # run CV decoder: change vs non-change
    model = {}
    rfc = RandomForestClassifier(class_weight='balanced')
    model['cv_prediction'] = cross_val_predict(rfc, X,y,cv=5)
    model['behavior_correlation'] = compute_behavior_correlation(\
        sample_cells[0].copy(),model)
    model['test_score'] = np.mean(y == model['cv_prediction'])

    # run CV decoder: hit vs miss
    X = []
    for cell in sample_cells:
        X.append(get_matrix(cell.query('is_change')))
    X = np.concatenate(X,axis=1)
    y = sample_cells[0].query('is_change')['hit'].astype(bool).values

    rfc = RandomForestClassifier(class_weight='balanced')
    model['cv_prediction_hit_vs_miss'] = cross_val_predict(rfc,X,y,cv=5)
    model['test_score_hit_vs_miss'] = np.mean(y == model['cv_prediction_hit_vs_miss'])

    # return results
    return model

def compute_behavior_correlation(cell, model):
    cell['prediction'] = model['cv_prediction']
    cell = cell.query('is_change').copy()
    cell['hit'] = cell['hit'].astype(bool)
    return cell[['prediction','hit']].corr()['hit']['prediction']

def plot_behavior_ratio(df,ncells=2,savefig=False,version=None):

    df = df.query('cre_line == "Slc17a7-IRES2-Cre"')
    df = df.query('experience_level == "Familiar"')
    df = df.query('n_cells > @ncells')

    x = df.groupby(['n_cells','ophys_experiment_id']).mean().reset_index()
    summary = x.groupby(['n_cells','visual_strategy_session'])['behavior_correlation'].mean()
    summary = summary.unstack()
    summary['ratio'] = summary[True]/summary[False]

    plt.figure(figsize=(5,4))
    plt.plot(summary.index.values, summary.ratio.values, 'ko-')
    plt.xlabel('# of cells',fontsize=16)
    plt.ylabel('correlation with behavior \n (ratio of visual/timing)',fontsize=16)
    ax = plt.gca()
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.xaxis.set_tick_params(labelsize=12)
    ax.yaxis.set_tick_params(labelsize=12)
    ax.set_ylim(0,4)
    ax.set_xlim(0,85)
    plt.tight_layout()

    # save decoder performance figure
    if savefig:
        filename='/allen/programs/braintv/workgroups/nc-ophys/alex.piet'+\
            '/behavior/decoding/figures_fit_{}/'.format(version)
        filename += 'exc_decoder_correlation_ratio.png'
        print(filename)
        plt.savefig(filename)


#!/usr/bin/env python3
import sys
import subprocess

from collections import namedtuple

size = 1024
# cpuiter = range(1, 56, 4)
# cpuiter = range(8,9)
# cpuiter = range(56,57)
# cpuiter = range(41,42)
# cpuiter = range(1, 4)
# cpuiter = range(5, 6)
cpuiter = range(15, 16)
# cpuiter = range(16, 17)

config_temp = namedtuple("Config", ["prgname", "verboselevel", "nx", "ny", "niters", "cpuiter"])
config =                config_temp("mpi no stencil", 2, size, size, 100, cpuiter)

## mpirun options:
##        -l or -prepend-rank
##              Use this option to insert the MPI process rank at the beginning of all lines written to the standard output.
##
##       -tune [<arg>]              Use this option to optimize the Intel MPI Library performance by using the data collected by the mpitune utility.
##
##       -trace <profiling library> or -t <profiling library>
##              Use  this  option  to  profile your MPI application with Intel(R) Trace Collector using the indicated profiling_library If you do not specify profiling_library , the default
##              profiling library libVT.so is used.
##
##       -trace-imbalance
##              Use this option to profile your MPI application with Intel Trace Collector using the libVTim.so library.
##
##       -mps   Use this option to collect statistics from your MPI application using MPI Performance Snapshot, which can collect hardware  performance  metrics,  memory  consumption  data,
##              internal MPI imbalance and OpenMP* imbalance statistics. When you use this option, a new folder stat_<date>-<time> is generated that contains the statistics. You can analyze
##              the collected data with the mps utility, which provides a high-level overview of your application performance.  MPI Performance Snapshot is available  as  part  of  Intel(R)
##              Trace Analyzer and Collector.
##
##       -check-mpi [<checking_library>]
##              Use this option to check your MPI application for correctness using the specified checking_library If you do not specify checking_library , the default checking library lib-
##              VTmc.so is used.
##
##       -trace-pt2pt
##              Use this option to collect the information about point-to-point operations using Intel Trace Analyzer and Collector. The option requires that you also use the -trace option.
##
##       -trace-collectives
##              Use this option

Job = namedtuple("Job", ["number_of_processes", "number_of_threads" ,"command", "args"])

def verify(ref_name):
    with open(ref_name, "rb") as stencil_ref:
        try:
            with open("stencil.pgm", "rb") as stencil_res:
                # Check image header
                ref_format, ref_nx, ref_ny, ref_depth = stencil_ref.readline().split()
                format, nx, ny, depth = stencil_res.readline().split()
                if not (format == ref_format == b"P5"):
                    print("Error: incorrect file format"); sys.exit()
                if not (depth == ref_depth == b"255"):
                    print("Error: incorrect depth value"); sys.exit()
                if not (nx == ref_nx and ny == ref_ny):
                    print("Error: image sizes do not match"); sys.exit()

                # Compare images
                passed = True
                for j in range(int(ref_ny)):
                    for i in range(int(ref_nx)):
                        ref_val = ord(stencil_ref.read(1))
                        val     = ord(stencil_res.read(1))
                        if abs(ref_val - val) > 1:
                            passed = False

                return passed
        except:
            return False

def run_job(job):
    if config.verboselevel:
        print("mpirun -genv OMP_NUM_THREADS {} -np {} {} {}"\
                .format(job.number_of_threads, job.number_of_processes, job.command, " ".join(job.args)))
    comp = subprocess.run(["mpirun",
                            "-genv", "OMP_NUM_THREADS", str(job.number_of_threads),
                           "-np",  str(job.number_of_processes),
                            job.command, *job.args,
                          ],
                          stdout=subprocess.PIPE,
                          stderr=subprocess.PIPE)
    return comp.stdout, comp.stderr

def get_runtime(outtext):
    if type(outtext)==type(b""):
        outtext = outtext.decode("latin1")
    i1 = outtext.find("runtime: ") + len("runtime: ")
    i2 = outtext.find(" s", i1)
    try:
        return float(outtext[i1:i2])
    except:
        return -1.0

def test_cpus_range(outname, nx, ny, niters, cpuiter, ver=False, vername=""):
    threaditer=range(1,2,1)
    print(nx, ny, niters)
    print("-"*20)
    false_ones = []
    for ncpus in cpuiter:
        for nthreads in threaditer:
            if ver:
                subprocess.run(["/bin/rm",  "stencil.pgm"])
            job = Job(ncpus, nthreads, outname, [str(nx), str(ny), str(niters)])
            stdout, stderr = run_job(job)
            print(ncpus, nthreads, get_runtime(stdout))
            if ver:
                f = verify(vername)
                print("Verify with", vername, ":", f)
                if not f:
                    false_ones.append(ncpus)
            if config.verboselevel == 2:
                print(stdout.decode("utf-8"))
                print(stderr.decode("utf-8"))
                print("="*20+2*'\n')
    if ver:
        print(false_ones)

print("hello from python")
print(config.prgname)
test_cpus_range("./out/mpistencil", config.nx, config.ny , config.niters, config.cpuiter, True, "clean.pgm")


# -*- coding: utf-8 -*-
"""
Created on Wed Nov 15 13:24:31 2023

@author: Diego
"""

import os
import string
import numpy as np
import pandas as pd
import datetime as dt

class DateGenerator:

    def _gen_rand_holidays(self, df: pd.DataFrame) -> pd.DataFrame:

        df_shuffled = df.sample(frac = 1)
        df_market_open = df_shuffled.head(250).assign(market_day = "open")
        good_dates = df_market_open.date.drop_duplicates().to_list()
        df_market_holiday = df_shuffled.query("date != @good_dates").assign(market_day = "holiday")

        df_out = pd.concat([df_market_open, df_market_holiday])
        return df_out

    # function to check that we have 250 trading days
    def _check_days_count(self):

        bad_data = (self.df_market.query(
            "market_day == 'open'")
            [["contract_name", "local_time"]].
            assign(
                date = lambda x: x.local_time.dt.strftime("%Y-%m-%d"),
                year = lambda x: x.local_time.dt.year).
            drop(columns = ["local_time"]).
            drop_duplicates().
            groupby(["contract_name", "year"]).
            agg("count")
            ["date"].
            reset_index().
            query("date != 250"))

        if len(bad_data) == 0:
            if self.verbose == True: print("Data has correct days per year")
        else:
            if self.verbose == True: print("Data does not have correct days per year")

    def _check_hours_count(self):

        # due to using year end and international times there are some days
        # with less than a full trading day which needs to be accounted for
        first_date = self.df_market.local_time.min().date().strftime("%Y-%m-%d")

        bad_data = (self.df_market.query(
            "market_hour == 'open'")
            [["contract_name", "local_time"]].
            assign(
                date_hour = lambda x: x.local_time.dt.strftime("%Y-%m-%d %H"),
                date = lambda x: x.local_time.dt.strftime("%Y-%m-%d")).
            drop(columns = ["local_time"]).
            drop_duplicates().
            groupby(["contract_name", "date"]).
            agg("count")
            ["date_hour"].
            reset_index().
            query("date_hour != 20 & date > @first_date"))

        if len(bad_data) == 0:
            if self.verbose == True: print("Data has correct hours per day")

        else:
            if self.verbose == True: print("Data does not have correct hours per day")

    # generates market trading days / hours / timezones
    def __init__(
            self,
            country_contract: dict = None,
            end_date: dt.datetime = dt.datetime(year = dt.date.today().year, month = 1, day = 1),
            # default first of current year
            year_lookback: int = 10,
            verbose: bool = True):

        self.seed = 1234
        np.random.seed(self.seed)

        self.verbose = verbose

        self.country_contract = {
            "NYC": 1,
            "Chicago": 1,
            "London": 1,
            "Tokyo": 1,
            "Frankfurt": 1}

        # when contract gets passed through just ensure it meets correct formatting
        if country_contract != None:

            for country, num in country_contract.items():

                if country not in list(self.country_contract.keys()):
                    raise ValueError("Only accepting NYC, Chicago, London, Tokyo, or Frankfurt contracts")

                if type(num) != int:
                    raise TypeError("Only accepting int for contract count")

            self.country_contract = country_contract

        self.contract_count = sum(list(self.country_contract.values()))

        # ensure year passed through is type int
        if type(year_lookback) != int: raise TypeError("year_lookback must be type int")
        if year_lookback < 2: raise ValueError("year_lookback must be greater than a year")

        self.end_date = end_date
        self.start_date = dt.datetime(
            year = self.end_date.year - year_lookback,
            month = self.end_date.month,
            day = self.end_date.day)

        if self.verbose == True:
            print("Generating {} contracts\n{} year lookback\nstart date: {}\nend date: {}".format(
                self.country_contract,
                year_lookback,
                self.start_date,
                self.end_date))

        self.dates = [
            self.start_date + dt.timedelta(minutes=x) for x in range(
                0, int((self.end_date - self.start_date).total_seconds() / 60) + 1, 5)]

        # define contract names
        self.contract_names = []
        for country in self.country_contract.keys():

            for i in range(self.country_contract[country]):

                contract_name = "{}{}".format(country, i+1)
                self.contract_names.append(contract_name)

        translation_table = str.maketrans('', '', string.digits)

        if self.verbose == True: print("Standardizing Time to NYC")

        # initilizing longer data with datetimes
        self.df_time = (pd.DataFrame(
            columns = self.contract_names,
            index = self.dates).
            reset_index().
            melt(id_vars = "index", var_name = "contract_name").
            drop(columns = ["value"]).
            rename(columns = {"index": "utc_time"}).
            assign(
                zone = lambda x: x.contract_name.str.translate(translation_table),
                utc_time = lambda x: x.utc_time.dt.tz_localize("UTC"),
                nyc_time = lambda x: x.utc_time.dt.tz_convert("America/New_York").dt.strftime("%Y-%m-%d %H:%M")).
            groupby(["contract_name", "nyc_time"]).
            head(1))

        if self.verbose == True: print("Adding Time Zone changes")

        # making dataframe of timezones with daylight savings to merge
        # unfortunately when using tz_convert it makes a type specific to zone so the melt changes it to class object
        # instead coerce datetimes to string and then coerce back to datetimes when joined
        self.df_timezones_add = (pd.DataFrame({
            "utc_time": self.df_time.utc_time.drop_duplicates().sort_values().to_list()}).
            assign(
                NYC = lambda x: x.utc_time.dt.tz_convert("America/New_York").dt.strftime("%Y-%m-%d %H:%M"),
                Chicago = lambda x: x.utc_time.dt.tz_convert("America/Chicago").dt.strftime("%Y-%m-%d %H:%M"),
                London = lambda x: x.utc_time.dt.tz_convert("Europe/London").dt.strftime("%Y-%m-%d %H:%M"),
                Tokyo = lambda x: x.utc_time.dt.tz_convert("Asia/Tokyo").dt.strftime("%Y-%m-%d %H:%M"),
                Frankfurt = lambda x: x.utc_time.dt.tz_convert("Europe/Berlin").dt.strftime("%Y-%m-%d %H:%M")).
            melt(id_vars = "utc_time", var_name = "zone", value_name = "local_time"))

        # tz_convert will occassionaly bring in some date times from previous years when switching UTC to a
        # specific datetime
        min_year = self.start_date.year
        max_year = self.end_date.year

        # merge time zones to get local times check for weekends and add in holidays
        self.df_date_combined = (self.df_time.merge(
            right = self.df_timezones_add, how = "inner", on = ["utc_time", "zone"]).
            drop(columns = ["utc_time"]).
            assign(
                nyc_time = lambda x: pd.to_datetime(x.nyc_time),
                local_time = lambda x: pd.to_datetime(x.local_time),
                weekday = lambda x: x.local_time.dt.weekday,
                date = lambda x: pd.to_datetime(x.local_time.dt.strftime("%Y-%m-%d")),
                year = lambda x: x.date.dt.year).
            query("year >= @min_year & year < @max_year").
            drop(columns = ["year"]))

        if self.verbose == True: print("Adding Holidays")

        # create mask-like dataframe to join back with holiday and open market days
        self.df_holiday = (self.df_date_combined[
            ["zone", "date"]].
            drop_duplicates().
            assign(
                year = lambda x: x.date.dt.year,
                weekday = lambda x: x.date.dt.weekday).
            query("weekday != [5,6]").
            groupby(["zone", "year"]).
            apply(self._gen_rand_holidays).
            reset_index(drop = True).
            sort_values("date").
            drop(columns = ["year"]))

        # join back holidays on date and fillna with closed since its a weekend
        self.df_open = (self.df_date_combined.merge(
            right = self.df_holiday,
            how = "outer",
            on = ["date", "weekday", "zone"]).
            assign(
                market_day = lambda x: x.market_day.fillna("closed"),
                hour = lambda x: x.local_time.dt.hour))

        if self.verbose == True: print("Adding Market Hours")

        # put market hours in
        # preferably use assign but since _ror is not accepted with np.where
        # slice dataframe by hour and concat
        closed_hours = [17, 18, 19, 20]

        # we can get our specific market hours
        self.df_hour_open = (self.df_open.query(
            "market_day == 'open' & hour != @closed_hours").
            assign(market_hour = "open"))

        # then join then back to the original data and fillna as closed
        self.df_market = (self.df_open.merge(
            right = self.df_hour_open,
            how = "outer",
            on = self.df_open.columns.to_list()).
            assign(market_hour = lambda x: x.market_hour.fillna("closed")))

        self._check_days_count()
        self._check_hours_count()

    def save_data(self):

        self.parent_path = os.path.abspath(os.path.join(os.getcwd(), os.pardir))
        self.data_path = os.path.join(self.parent_path, "data")

        if os.path.exists(self.data_path) == False: os.makedirs(self.data_path)
        self.file_out = os.path.join(self.data_path, "date.parquet")
        self.df_market.to_parquet(path = self.file_out, engine = "pyarrow")

        if self.verbose == True: print("File Written to", self.file_out)

if __name__ == "__main__":

    date_generator = DateGenerator()
    date_generator.save_data()


'''
字符串编码一直是令人非常头疼的问题，尤其是我们在处理一些不规范的第三方网页的时候。虽然Python提供了Unicode表示的str和bytes两种数据类型，并且可以通过encode()和decode()方法转换，但是，在不知道编码的情况下，对bytes做decode()不好做。

对于未知编码的bytes，要把它转换成str，需要先“猜测”编码。猜测的方式是先收集各种编码的特征字符，根据特征字符判断，就能有很大概率“猜对”。

当然，我们肯定不能从头自己写这个检测编码的功能，这样做费时费力。chardet这个第三方库正好就派上了用场。用它来检测编码，简单易用。
'''
# 使用chardet
import chardet
# 当我们拿到一个bytes时，就可以对其检测编码。用chardet检测编码，只需要一行代码：
a = chardet.detect(b'Hello, world!')

'''检测出的编码是ascii，注意到还有个confidence字段，表示检测的概率是1.0（即100%）。

我们来试试检测GBK编码的中文：'''
data = '离离原上草，一岁一枯荣'.encode('gbk')
chardet.detect(data)
# 检测的编码是GB2312，注意到GBK是GB2312的超集，两者是同一种编码，检测正确的概率是74%，language字段指出的语言是'Chinese'。
# 对UTF-8编码进行检测:
data = '离离原上草，一岁一枯荣'.encode('utf-8')
chardet.detect(data)
# 日语
data = '最新の主要ニュース'.encode('euc-jp')
chardet.detect(data)


import sqlite3
import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from sklearn.metrics import f1_score, confusion_matrix
from torch.utils.data import Dataset, DataLoader
from models import *
import torch
from sklearn.decomposition import PCA
from tqdm import tqdm
import sys
from matplotlib import pyplot as plt

connection = sqlite3.connect("./DB/drugs.db")
print("Loading Data...")
df = pd.read_sql_query("SELECT * from train_set0", connection)
# print(df.head())
print(df.shape)
# df["target"].value_counts()[1000:1100].plot(kind="bar")
unique_pbdid = set(",".join(df["target"].values.flatten().tolist()).split(","))
pbdid_dict = {v:k for k, v in dict(enumerate(unique_pbdid)).items()}
del unique_pbdid
print(len(pbdid_dict))
print("Data Loaded...")
df = df.sample(frac=1).reset_index(drop=True)
df.isna().sum()

x = df.iloc[:, :-1].values
y = df.iloc[:, -1].values
# del df
scaler = StandardScaler()
scaler.fit(x[:, 1:])

x_t, x_v, y_t, y_v = train_test_split(x, y, test_size=0.01, random_state=42)
x_t[:, 1:] = scaler.transform(x_t[:, 1:])
x_v[:, 1:] = scaler.transform(x_v[:, 1:])
del x
del y
print("---Features Scaled---")

class GraphDataset(Dataset):
    def __init__(self, x, y, pbdid_dict):
        self.cids = df["cid"].unique()
        self.x = x
        self.y = y
        self.pbdid_dict = pbdid_dict

    def __len__(self):
        return self.x.shape[0]

    def __getitem__(self, idx):
        cid = self.cids[idx]
        index = np.where(self.x[:, 0] == cid)[0]
        x = self.x[index, 1:].tolist()
        y = np.zeros(len(self.pbdid_dict))
        targets = ",".join(self.y[index]).split(",")
        targets = [self.pbdid_dict[_id] for _id in targets if _id != ""]
        y[targets] = 1
        return torch.tensor(x), torch.tensor(y.tolist())

train_data = GraphDataset(x_t, y_t, pbdid_dict)
val_data = GraphDataset(x_v, y_v, pbdid_dict)
#
train_loader = DataLoader(train_data, shuffle=True, batch_size=64)
val_loader = DataLoader(val_data, shuffle=True, batch_size=64)

# x, y = next(iter(train_loader))
# print(x.shape, y.shape)
del df
# del x_t
# del y_t
# del x_v
# del y_v
# x, y = next(iter(train_loader))
# print(x, y)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = LinearClassifier(len(pbdid_dict)).to(device)
# criterion = torch.nn.NLLLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)

train=False
if train:
    if len(sys.argv) > 1:
        EPOCHS = int(sys.argv[1])
    else:
        EPOCHS = 10

    print("epoch\t train_loss\t val_loss")
    for epoch in range(1, EPOCHS+1):
        train_loss = 0.0
        val_loss = 0.0
        train_loop = tqdm(train_loader, total=len(train_loader), leave=False)
        val_loop = tqdm(val_loader, total=len(val_loader), leave=False)
        model.train()
        for batch in train_loop:
            x, y = batch
            x, y = x.to(device), y.to(device)
            optimizer.zero_grad()
            yhat, loss = model(x.float(), y)
            # loss = criterion(yhat, y)
            train_loss += loss.item()
            loss.backward()
            optimizer.step()

            train_loop.set_description(f"Epoch [{epoch}/{EPOCHS}]")
            train_loop.set_postfix(loss=loss.item())

        model.eval()
        with torch.no_grad():
            for batch in val_loop:
                x, y = batch
                x, y = x.to(device), y.to(device)
                yhat, loss = model(x.float(), y)
                # loss = criterion(yhat, y)
                val_loss += loss.item()
        print(f"{epoch}\t {round(train_loss/len(train_loader), 4)}\t {round(val_loss/len(val_loader), 4)}")
        torch.save(model.state_dict(), f"./models/lc-stage-2-{epoch}.pt")
valid = False
if valid:
    model.load_state_dict(torch.load("./models/lc-stage-2-10.pt", map_location=torch.device("cpu")))
    val_loop = tqdm(val_loader, total=len(val_loader), leave=False)
    correct = 0
    total = 0
    model.eval()
    with torch.no_grad():
        for batch in val_loop:
            x, y = batch
            x, y = x.to(device), y.to(device)
            yhat, _ = model(x.float(), y)
            yhat = yhat.argmax(1)
            predictions = yhat == y
            correct += sum(predictions.numpy().astype(int))
            total += y.shape[0]
    print("Model Accuracy", round(correct/total, 4) * 100.0, "%")


import uuid

from computedfields.models import ComputedFieldsModel, computed
from django.core import validators
from django.db import models

from categories import models as categories_models

class Product(ComputedFieldsModel):
    id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
    name = models.CharField(max_length=100)
    description = models.TextField()
    color = models.CharField(max_length=100)
    price = models.DecimalField(
        max_digits=9, decimal_places=2, validators=[validators.MinValueValidator(0)]
    )
    category = models.ForeignKey(categories_models.Category, on_delete=models.CASCADE)

    @computed(
        models.DecimalField(
            help_text="Promotion price based on the product category discount",
            validators=[validators.MinValueValidator(0)],
            max_digits=9,
            decimal_places=2,
            default=0,
            null=False,
        ),
        depends=[("self", ["price"]), ("category", ["discount"])],
    )
    def promotion_price(self) -> float:
        return round(float(self.price) * float(1 - self.category.discount_for_calc), 2)

    class Meta:
        indexes = [
            models.Index(
                fields=["name", "description"],
                name="name_description_filter_index",
            ),
        ]