# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/00_Create Project.ipynb (unless otherwise specified).

__all__ = ['init_available_projects', 'summarize_h5', 'pgrid_to_run_parameters', 'collect_statistics', 'fix_geojson',
           'main']

# Cell
from pathlib import Path
import shutil
import os
import pandas as pd
import geopandas as gpd
import shapely
from shapely.geometry.polygon import Polygon
from shapely.ops import cascaded_union, unary_union
import numpy as np
from time import time
from typing import *
import junevis.path_fixes as pf
import json

import junevis.process_loggers as process_loggers

# Cell
def init_available_projects(project_name: str):
    pf.AVAILABLE_PROJECTS.touch()

    with open(str(pf.AVAILABLE_PROJECTS), 'r+') as fp:
        available_projects = set([p.strip() for p in fp.readlines()])
        if project_name in available_projects:
            if not force_add_project:
                raise ValueError(f"Cannot create project of name '{project_name}': Project already exists in {pf.AVAILABLE_PROJECTS}"
    )
            else:
                shutil.rmtree(outdir) # Delete existing project of that name
                fp.truncate(0); fp.seek(0); # Delete file contents
                available_projects.remove(project_name)
        return available_projects

# Cell
def summarize_h5(record_f, outdir):
    """Dependent on the context variable `output_dir`. The actual summarized output is much smaller than the record file itself"""
    start = time()
    runId = record_f.stem.split("_")[1]
    print(f"Processing {runId}: ")
    df = process_loggers.regional_outputs(record_f)

    # Add cumulative columns
    region_grouped_df = df.groupby(level=0)
    df['currently_dead'] = region_grouped_df.deaths.cumsum()
    df['currently_recovered'] = region_grouped_df.recovered.cumsum()

    # Rename region
    df = df.rename_axis(index=["region", "timestamp"])

    outfile = outdir / f"summary_{int(runId):03}.csv"
    print(f"Saving to {str(outfile)}")
    df.to_csv(str(outfile))
    print(f"\nTook {time() - start} seconds")
    print("\n-------\n")
    return df

# Cell
def pgrid_to_run_parameters(parameter_grid: dict) -> dict:
    """Convert parameter_grid dictionary to desired metadata dictionary"""
    run_parameters = {}

    # Create run_parameters
    for k, v in parameter_grid.items():
        for i in range(len(v)):
            curr = run_parameters.get(str(i), {})
            curr[k] = v[i]
            run_parameters[str(i)] = curr

    params_varied = list(parameter_grid.keys())

    return {
        "parameters_varied": params_varied,
        "run_parameters": run_parameters,
    }

# Cell
def collect_statistics(project: Union[str, Path]):
    project = Path(project)
    csvfs = list(project.glob("summary*.csv"))
    dfs = [pd.read_csv(csvf) for csvf in csvfs]
    big_df = pd.concat(dfs, ignore_index=True)

    all_regions = list(set(big_df.region))
    all_fields = [f for f in big_df.columns if f != "Unnamed: 0"]
    all_timestamps = list(set(big_df.timestamp))

    string_fields = ["timestamp", "region", "Unnamed: 0"]
    numerical_fields = [f for f in all_fields if f not in string_fields]
    big_df_num = big_df.loc[:, numerical_fields]
    max_vals = big_df_num.max(axis=0)
    min_vals = big_df_num.min(axis=0)

    df_minmax = pd.DataFrame([max_vals, min_vals], index=["max", "min"])
    field_minmaxes = df_minmax.to_dict(orient="dict")

    return {
        "all_regions": sorted(all_regions),
        "all_timestamps": sorted(all_timestamps),
        "all_fields": sorted(all_fields),
        "field_statistics": field_minmaxes
    }

# Cell
def fix_geojson(gjson_file):
    gdf = gpd.read_file(gjson_file)

    for i, shape in enumerate(gdf.geometry):
        # To reduce the shape of the multipolygon, take the shape of the largest area
        if shape.geom_type == "MultiPolygon":
            polygon = shape[np.argmax(np.array([p.area for p in shape]))]
        else:
            polygon = shape
        gdf.geometry[i] = polygon

    # The frontend operates with a `SSID` field instead of a `region` field to name each area.
    gdf['SSID'] = gdf['region']
    return gdf

# Cell
from fastcore.script import *

@call_parse
def main(record_path:Param("Path to JUNE simulation records and parameter grid", str),
         force_add_project:Param("Overwrite project if it already exists", store_true)=False,
         test_only:Param("Test behavior without changing files", store_true)=False,
         project_name:Param("Name the project. If not provided, use folder name of `record_path`", str)=None,
         description:Param("Description of project", str)="NA",
        ):
    """Create a project that can be visualized from the record files"""

    base = Path(record_path) # Path where loggers and parameter grid are stored
    project_name = base.stem if project_name is None else project_name
    output_dir = pf.PROJECTS / project_name
    if not output_dir.exists() and not test_only: output_dir.mkdir(parents=True)

    active_projects = init_available_projects(project_name)

    record_names = sorted(list(base.glob("*.h5")))
    for r in record_names:
        print(f"Summarizing {r}")
        if not test_only: df = summarize_h5(r, output_dir)

    print("ALL SUMMARIES COMPLETED\n-------------\n-------------\n")

    # Once the summary files have been created, we can accumulate the statistics into the `metadata.json` file
    print("Creating metadata...")
    with open(base / "parameter_grid.json") as fp:
        parameter_grid = json.load(fp)
    param_info = pgrid_to_run_parameters(parameter_grid)
    project_stats = collect_statistics(output_dir)

    # Now we can save the metadata for this project, including the optional description
    metadata = {"description": description}; [metadata.update(p) for p in [param_info, project_stats]];
    if not test_only:
        with open(output_dir / "metadata.json", 'w+') as fp:
            json.dump(metadata, fp, indent=4)

    # Copy over the geography description
    print("Fixing geojson...")
    gdf = fix_geojson(base / "sites.geojson")
    if not test_only: gdf.to_file(output_dir / "sites.new.geojson", driver='GeoJSON')

    # Add to available projects
    print(f"Adding '{project_name}' to {pf.AVAILABLE_PROJECTS}")
    new_available_projects = "".join(["\n" + p for p in (list(active_projects) + [project_name])]).strip()
    print(f"New projects: {new_available_projects}")

    if not test_only:
        with open(pf.AVAILABLE_PROJECTS, 'r+') as fp:
            fp.write(new_available_projects)

    print("COMPLETE")