#!/usr/bin/env python3
"""Local restricted-torque migration-time map, standard library only.

Fixed mass4 Earth, radius1 AU, height0.055, density1e-4, eccentricity0.
This is not an evolving-fluid simulation, posterior or observed migration map.
"""
from pathlib import Path
import argparse, json, math, html
from disk_torques import (
    DEFAULT_DISK,
    EARTH_SOLAR_MASS_RATIO,
    YEARS_PER_CODE_TIME,
    torque_components,
)


def calculate():
    cells = []
    for slope_index in range(25):
        slope = 0.3 + slope_index * 0.05
        for viscosity_index in range(31):
            exponent = -9 + viscosity_index / 6
            disk = {
                **DEFAULT_DISK,
                "surfaceDensityAtOneAu": 1e-4,
                "temperatureSlope": slope,
                "viscosity": 10**exponent,
            }
            mass = 4 * EARTH_SOLAR_MASS_RATIO
            torque = torque_components(mass, 1, 0, disk)
            drift = 2 * torque["total"] / mass / YEARS_PER_CODE_TIME
            time = None if drift == 0 else 1 / abs(drift)
            cells.append(
                {
                    "temperatureSlope": slope,
                    "viscosityExponent": exponent,
                    "viscosity": 10**exponent,
                    "waveCoefficient": torque["wave"] / torque["normalization"],
                    "corotationCoefficient": torque["corotation"]
                    / torque["normalization"],
                    "totalCoefficient": torque["total"] / torque["normalization"],
                    "driftAuPerJulianYear": drift,
                    "localMigrationTimeJulianYears": time,
                    "signStableUnderComponentSensitivity": abs(torque["total"])
                    > 0.2 * (abs(torque["wave"]) + abs(torque["corotation"])),
                }
            )
    return {
        "kind": "Synthetic local migration-time map from the restricted Paardekooper prescription, not hydrodynamics or observations",
        "fixed": {
            "outerMassEarth": 4,
            "radiusAu": 1,
            "eccentricity": 0,
            "aspectRatio": 0.055,
            "surfaceDensityCodeUnits": 1e-4,
            "densitySlope": 0.5,
            "physicalGamma": 1.4,
            "prandtl": 0.5,
            "fitSofteningOverHeight": 0.4,
        },
        "axes": {"viscosityExponent": [-9, -4, 31], "temperatureSlope": [0.3, 1.5, 25]},
        "timeDefinition": "Local radius divided by absolute circular drift rate; null is infinite at exact zero torque. The local time is not an integrated arrival time.",
        "sensitivityMeaning": "A deterministic plus/minus20 percent perturbation envelope on each torque contribution; illustrative sensitivity, not a statistical confidence or calibrated model-error interval. False means net sign can reverse within that envelope.",
        "cells": cells,
    }


def figure(data):
    width, height = 340, 410
    left, right, top, bottom = 53, 326, 48, 276
    parts = [
        f'<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}" viewBox="0 0 {width} {height}" role="img"><title>Local synthetic migration-time map</title><rect width="100%" height="100%" fill="#091226"/>'
    ]

    def text(x, y, value, size=12, anchor="start", colour="#cbd5e1"):
        parts.append(
            f'<text x="{x:.3f}" y="{y:.3f}" font-family="system-ui,sans-serif" font-size="{size}" text-anchor="{anchor}" fill="{colour}">{html.escape(str(value))}</text>'
        )

    text(left, 20, "Local migration time", 16)
    text(left, 38, "Temperature slope", 12)
    cell_width = (right - left) / 31
    cell_height = (bottom - top) / 25
    for index, cell in enumerate(data["cells"]):
        row, column = divmod(index, 31)
        x = left + column * cell_width
        y = bottom - (row + 1) * cell_height
        timescale = cell["localMigrationTimeJulianYears"] or 1e20
        strength = max(0, min(1, (7 - math.log10(timescale)) / 3))
        base = (85, 210, 241) if cell["driftAuPerJulianYear"] > 0 else (249, 161, 89)
        colour = "#" + "".join(
            f"{round(13+(channel-13)*strength):02x}" for channel in base
        )
        parts.append(
            f'<rect x="{x:.3f}" y="{y:.3f}" width="{cell_width+.1:.3f}" height="{cell_height+.1:.3f}" fill="{colour}"/>'
        )
        if not cell["signStableUnderComponentSensitivity"]:
            parts.append(
                f'<circle cx="{x+cell_width/2:.3f}" cy="{y+cell_height/2:.3f}" r=".65" fill="#e2e8f0"/>'
            )
    for tick in range(6):
        text(
            left + (right - left) * tick / 5,
            bottom + 20,
            str(tick - 9),
            anchor="middle",
        )
    for tick in range(4):
        text(
            left - 8,
            bottom - (bottom - top) * tick / 3 + 4,
            f"{.3+1.2*tick/3:.1f}",
            anchor="end",
        )
    text(left, 314, "Viscosity exponent", 12)
    text(left, 333, "Orange: inward · cyan: outward", 11)
    for index, timescale in enumerate([1e4, 1e5, 1e6, 1e7]):
        x = left + index * (right - left) / 4
        strength = max(0, min(1, (7 - math.log10(timescale)) / 3))
        for sign, base in enumerate([(249, 161, 89), (85, 210, 241)]):
            colour = "#" + "".join(
                f"{round(13+(channel-13)*strength):02x}" for channel in base
            )
            parts.append(
                f'<rect x="{x:.3f}" y="{347+sign*7}" width="42" height="6" fill="{colour}"/>'
            )
        text(x + 21, 376, ["10k", "100k", "1M", "10M"][index], 11, "middle")
    text(left, 399, "Time in years · dots: sign sensitivity", 11)
    return "".join(parts) + "</svg>\n"


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--output", type=Path, default=Path("migration-timescale-map.json")
    )
    parser.add_argument(
        "--figure", type=Path, default=Path("migration-timescale-map.svg")
    )
    args = parser.parse_args()
    data = calculate()
    args.output.write_text(json.dumps(data, indent=2, allow_nan=False) + "\n")
    args.figure.write_text(figure(data))
    print(
        json.dumps(
            {
                "cells": len(data["cells"]),
                "inward": sum(
                    cell["driftAuPerJulianYear"] < 0 for cell in data["cells"]
                ),
                "outward": sum(
                    cell["driftAuPerJulianYear"] > 0 for cell in data["cells"]
                ),
                "signSensitive": sum(
                    not cell["signStableUnderComponentSensitivity"]
                    for cell in data["cells"]
                ),
            }
        )
    )


if __name__ == "__main__":
    main()