"""Create the eight lightweight entry notebooks for the organized project."""

from pathlib import Path
import json

ROOT = Path(__file__).resolve().parent


def notebook(title, imports, analysis):
    def cell(kind, source):
        block = {"cell_type": kind, "metadata": {}, "source": source.splitlines(True)}
        if kind == "code":
            block.update({"execution_count": None, "outputs": []})
        return block

    setup = """from pathlib import Path
import sys

PROJECT_ROOT = Path.cwd().resolve()
while not (PROJECT_ROOT / 'data.py').exists():
    if PROJECT_ROOT.parent == PROJECT_ROOT:
        raise FileNotFoundError('Open this notebook from inside organized_uncertainty_analysis.')
    PROJECT_ROOT = PROJECT_ROOT.parent
if str(PROJECT_ROOT.parent) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT.parent))
"""
    return {
        "cells": [
            cell("markdown", f"# {title}\n"),
            cell("code", setup + "\n" + imports),
            cell("code", analysis),
        ],
        "metadata": {
            "kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},
            "language_info": {"name": "python", "version": "3"},
        },
        "nbformat": 4,
        "nbformat_minor": 5,
    }


specs = {
    "logistic/notebooks/01_data_fit_gof.ipynb": (
        "Logistic regression: data, fitting, and goodness of fit",
        "from organized_uncertainty_analysis.logistic import data_fit_gof as fit\n",
        "import pandas as pd\ndata = fit.prepare_full_trim(PROJECT_ROOT)\nrows = []\nfor name, (x, y) in data.items():\n    for transform in ('raw', 'log'):\n        pack = fit.fit_logistic_x(x, y, transform=transform, l2=1e-8)\n        x_model = x if transform == 'raw' else __import__('numpy').log(x)\n        rows.append({'Dataset': name, 'Predictor': transform.upper(), **fit.goodness_of_fit(x_model, y, pack, l2=1e-8)})\ngof_table = pd.DataFrame(rows)\ngof_table.to_csv(PROJECT_ROOT / 'outputs' / 'logistic_goodness_of_fit.csv', index=False)\ngof_table\n",
    ),
    "logistic/notebooks/02_ci_estimation.ipynb": (
        "Logistic regression: confidence intervals and bands",
        "from organized_uncertainty_analysis.logistic.data_fit_gof import prepare_full_trim\nfrom organized_uncertainty_analysis.logistic import ci_estimation as ci\n",
        "B = 5000\ndata = prepare_full_trim(PROJECT_ROOT)\nP = {}\nfor dataset, (x, y) in data.items():\n    for transform in ('raw', 'log'):\n        key = f'{dataset}-{transform.upper()}'\n        P[key] = ci.fit_ci_pack_rawgrid(x, y, transform=transform, name=key, l2=1e-8, B=B, seed=123)\nsummary = ci.param_ci_table_4methods(P)\nsummary.to_csv(PROJECT_ROOT / 'outputs' / 'logistic_parameter_ci.csv', index=False)\nfig, axes = ci.plot_ci_four_panels(P)\nfig.savefig(PROJECT_ROOT / 'outputs' / 'logistic_ci_bands.png', dpi=300, bbox_inches='tight')\nfig.savefig(PROJECT_ROOT / 'outputs' / 'logistic_ci_bands.pdf', bbox_inches='tight')\nsummary\n",
    ),
    "logistic/notebooks/03_elasticity.ipynb": (
        "Logistic regression: elasticity",
        "import numpy as np\nfrom organized_uncertainty_analysis.logistic.data_fit_gof import prepare_full_trim, fit_logistic_x\nfrom organized_uncertainty_analysis.logistic import elasticity as el\n",
        "data = prepare_full_trim(PROJECT_ROOT)\nP = {}\nfor dataset, (x, y) in data.items():\n    for transform in ('raw', 'log'):\n        key = f'{dataset}-{transform.upper()}'\n        P[key] = {'b': fit_logistic_x(x, y, transform=transform, l2=1e-8), 'transform': transform}\nelasticity_table = el.elasticity_table_4panels(P, make_plots=False)\nelasticity_table.to_csv(PROJECT_ROOT / 'outputs' / 'logistic_elasticity.csv', index=False)\nfig, axes = el.plot_combined_elasticity(elasticity_table)\nfig.savefig(PROJECT_ROOT / 'outputs' / 'logistic_elasticity.png', dpi=600, bbox_inches='tight')\nfig.savefig(PROJECT_ROOT / 'outputs' / 'logistic_elasticity.pdf', bbox_inches='tight')\nelasticity_table\n",
    ),
    "logistic/notebooks/04_noise_measurement.ipynb": (
        "Logistic regression: measurement noise",
        "from organized_uncertainty_analysis.logistic.data_fit_gof import prepare_full_trim\nfrom organized_uncertainty_analysis.logistic import noise_measurement as noise\n",
        "data = prepare_full_trim(PROJECT_ROOT)\nfig, axes, full_noise, trim_noise = noise.make_noise_figure(*data['FULL'], *data['TRIM'], transform='raw', n_refit=1100, n_tta=10000, save_path=PROJECT_ROOT / 'outputs' / 'logistic_noise')\n",
    ),
    "bayesian/notebooks/01_data_fit_gof.ipynb": (
        "Constrained Bayesian model: data, fitting, and goodness of fit",
        "import os\nos.chdir(PROJECT_ROOT)\nfrom organized_uncertainty_analysis.bayesian import data_fit_gof as fit\n",
        "import pandas as pd\nfit_result = fit.fit_full_trim(seed=123)\ngof_table = pd.DataFrame([{'Dataset': name, **fit_result[name]['gof']} for name in ('FULL', 'TRIM')])\ngof_table.to_csv(PROJECT_ROOT / 'outputs' / 'bayesian_goodness_of_fit.csv', index=False)\ngof_table\n",
    ),
    "bayesian/notebooks/02_ci_estimation.ipynb": (
        "Constrained Bayesian model: confidence intervals and bands",
        "import os\nos.chdir(PROJECT_ROOT)\nfrom organized_uncertainty_analysis.bayesian import ci_estimation as ci\n",
        "ci_result = ci.run_bayesian_ci(B_nonpar=400, B_param=400, M_mca=10000, max_hessian_condition=1e6)\nparameter_table = ci.make_parameter_ci_table(ci_result)\nderived_table = ci.make_derived_ci_table(ci_result)\nparameter_table.to_csv(PROJECT_ROOT / 'outputs' / 'bayesian_parameter_ci.csv', index=False)\nderived_table.to_csv(PROJECT_ROOT / 'outputs' / 'bayesian_x50_s50_ci.csv', index=False)\nfig, axes = ci.plot_bayesian_ci(ci_result, band_support='observed')\nfig.savefig(PROJECT_ROOT / 'outputs' / 'bayesian_ci_bands.png', dpi=300, bbox_inches='tight')\nfig.savefig(PROJECT_ROOT / 'outputs' / 'bayesian_ci_bands.pdf', bbox_inches='tight')\nparameter_table, derived_table\n",
    ),
    "bayesian/notebooks/03_elasticity.ipynb": (
        "Constrained Bayesian model: elasticity",
        "import os\nos.chdir(PROJECT_ROOT)\nfrom organized_uncertainty_analysis.bayesian.data_fit_gof import fit_full_trim\nfrom organized_uncertainty_analysis.bayesian import elasticity as el\n",
        "fit_result = fit_full_trim(seed=123)\nelasticity_table = el.elasticity_full_trim(fit_result)\nelasticity_table.to_csv(PROJECT_ROOT / 'outputs' / 'bayesian_elasticity.csv', index=False)\nfig, axes = el.plot_combined_elasticity(elasticity_table)\nfig.savefig(PROJECT_ROOT / 'outputs' / 'bayesian_elasticity.png', dpi=600, bbox_inches='tight')\nfig.savefig(PROJECT_ROOT / 'outputs' / 'bayesian_elasticity.pdf', bbox_inches='tight')\nelasticity_table\n",
    ),
    "bayesian/notebooks/04_noise_measurement.ipynb": (
        "Constrained Bayesian model: measurement noise",
        "import os\nos.chdir(PROJECT_ROOT)\nfrom organized_uncertainty_analysis.bayesian import noise_measurement as noise\n",
        "noise_result = noise.run_noise_analysis(n_refit=150, n_tta=300, seed=1234)\nsummary = noise.make_noise_summary_table(noise_result)\nfig, axes = noise.plot_4panels(noise_result['full'], noise_result['trim'], save_prefix=PROJECT_ROOT / 'outputs' / 'bayesian_noise')\nsummary\n",
    ),
}

for relative_path, args in specs.items():
    path = ROOT / relative_path
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(notebook(*args), indent=1), encoding="utf-8")
    print(path.relative_to(ROOT))