import numpy as np

from nuspacesim import eas_optical


def test_known_values():
    altitude = np.arange(0, 21, dtype=np.float32)
    beta_tr = np.full_like(altitude, np.radians(42.0), dtype=np.float32)
    cpa = eas_optical.eas.CphotAng()

    dphots, cang = cpa(beta_tr, altitude)

    # print(dphots)
    # print(cang)

    old_dphots = np.array(
        [
            1.6210271e03,
            1.8776989e03,
            2.1580850e03,
            2.4297334e03,
            2.6037090e03,
            2.5893950e03,
            2.6665376e03,
            1.9930570e03,
            1.2905229e03,
            7.2669135e02,
            3.5919400e02,
            1.5781609e02,
            6.3526283e01,
            2.4607666e01,
            9.4170065e00,
            3.5231948e00,
            1.3178004e00,
            5.0189447e-01,
            1.9809218e-01,
            8.2580306e-02,
            3.6842916e-02,
        ]
    )

    old_cang = np.array(
        [
            1.1297,
            1.0476834,
            0.9589647,
            0.86524945,
            0.7722871,
            0.68627226,
            0.6104225,
            0.54487425,
            0.48861432,
            0.4402858,
            0.39882657,
            0.3633839,
            0.33331406,
            0.30719605,
            0.28355855,
            0.2634559,
            0.24582747,
            0.22999273,
            0.21532202,
            0.20063952,
            0.18477382,
        ]
    )

    assert np.allclose(dphots, old_dphots)
    assert np.allclose(cang, old_cang)


if __name__ == "__main__":
    # test_known_values()

    import matplotlib.pyplot as plt

    plt.style.use("seaborn-whitegrid")
    plt.rcParams["font.size"] = 16

    degs = list([1.0, 5.0, 10.0, 20.0, 30.0])
    ds = []
    cs = []

    jds = [
        np.array(
            [
                5.27674332e-04,
                0.409867704,
                3.94313812,
                10.1881313,
                18.0240078,
                28.4657440,
                39.6969986,
                54.5465965,
                73.7063370,
                97.1183090,
                127.589645,
                174.656738,
                238.554108,
                326.405029,
                452.014099,
                642.060242,
                904.732117,
                1232.60498,
                1535.39136,
                2066.62646,
                1633.27539,
            ]
        ),
        np.array(
            [
                9.30504322,
                25.3226070,
                42.3239822,
                58.3105316,
                74.3881149,
                92.6042557,
                113.324104,
                139.367432,
                156.780945,
                180.987335,
                228.783035,
                297.147156,
                389.135193,
                517.803101,
                684.509583,
                866.451965,
                993.215820,
                1206.91724,
                901.117920,
                556.321472,
                290.776581,
            ]
        ),
        np.array(
            [
                63.6738396,
                96.9967651,
                124.916992,
                151.680115,
                180.303146,
                212.641663,
                251.588882,
                298.172302,
                359.177368,
                447.978363,
                574.827698,
                744.875977,
                928.580750,
                1055.05176,
                1268.87500,
                986.335632,
                643.666016,
                358.620361,
                175.488449,
                78.0030975,
                32.5095978,
            ]
        ),
        np.array(
            [
                319.989349,
                392.296692,
                463.493835,
                540.937988,
                631.059265,
                744.649292,
                896.025574,
                1093.49854,
                1304.78015,
                1446.11072,
                1504.65527,
                1362.30811,
                905.672241,
                510.569275,
                250.467560,
                110.092865,
                44.7798576,
                17.4202747,
                6.75650692,
                2.54880404,
                0.951708376,
            ]
        ),
        np.array(
            [
                791.055542,
                929.935425,
                1084.46570,
                1267.08813,
                1485.99060,
                1725.16589,
                1917.01294,
                1957.82263,
                2058.22607,
                1530.02161,
                965.538391,
                520.899170,
                246.437302,
                105.856003,
                42.3585625,
                16.4105015,
                6.22363472,
                2.31331372,
                0.857409120,
                0.323087066,
                0.125590757,
            ]
        ),
    ]

    jcs = [
        np.array(
            [
                1.32911205,
                1.27261925,
                1.21342850,
                1.15533781,
                1.09440362,
                1.03490627,
                0.975558519,
                0.914052486,
                0.852542639,
                0.792654872,
                0.732216477,
                0.670476615,
                0.611326814,
                0.555313587,
                0.499191791,
                0.443307370,
                0.391291529,
                0.341089129,
                0.293872952,
                0.252725810,
                0.219138503,
            ]
        ),
        np.array(
            [
                1.31018233,
                1.25294042,
                1.19424653,
                1.13343668,
                1.07256615,
                1.00931919,
                0.947368860,
                0.884270132,
                0.821494460,
                0.757842660,
                0.694497168,
                0.629486024,
                0.566345632,
                0.504391909,
                0.444903165,
                0.386975050,
                0.333529383,
                0.287159741,
                0.249708071,
                0.220059916,
                0.196107492,
            ]
        ),
        np.array(
            [
                1.28595519,
                1.22553182,
                1.16376925,
                1.10031724,
                1.03530788,
                0.969625294,
                0.902465105,
                0.833714664,
                0.764868081,
                0.695196211,
                0.623690188,
                0.549934447,
                0.478442669,
                0.413230062,
                0.356946319,
                0.310444713,
                0.273142755,
                0.242978141,
                0.218469203,
                0.198584095,
                0.182161078,
            ]
        ),
        np.array(
            [
                1.23594880,
                1.16961718,
                1.10034442,
                1.02877235,
                0.955798030,
                0.880288124,
                0.800689518,
                0.717127085,
                0.633447826,
                0.554169774,
                0.483273327,
                0.422520190,
                0.371766686,
                0.330366224,
                0.296454549,
                0.268502116,
                0.245628387,
                0.226084396,
                0.208619133,
                0.193726748,
                0.180828407,
            ]
        ),
        np.array(
            [
                1.18540561,
                1.11206508,
                1.03592265,
                0.956088781,
                0.870580733,
                0.780635655,
                0.692408621,
                0.610369325,
                0.538554668,
                0.476793230,
                0.424100250,
                0.379572451,
                0.342216313,
                0.311301470,
                0.285714626,
                0.263185740,
                0.243161157,
                0.226112887,
                0.211174414,
                0.197692946,
                0.184871987,
            ]
        ),
    ]

    alt = np.arange(0, 21, dtype=np.float32)

    for v in degs:
        beta_tr = np.full_like(alt, np.radians(v), dtype=np.float32)
        cpa = eas_optical.eas.CphotAng()

        dphots, cang = cpa(beta_tr, alt)
        ds.append(dphots)
        cs.append(cang)

    fig, axs = plt.subplots(2, 2, sharex=True, constrained_layout=True)
    colors = ["red", "orange", "green", "blue", "purple"]

    for i in range(5):
        axs[0, 0].plot(
            alt,
            ds[i],
            "--",
            label=f"AWR: {degs[i]}" + "\N{DEGREE SIGN}",
            color=colors[i],
        )
        axs[0, 0].plot(
            alt,
            jds[i],
            "-",
            label=f"JFK: {degs[i]}" + "\N{DEGREE SIGN}",
            color=colors[i],
        )
        axs[0, 0].set_title("Photon Density")
        axs[0, 0].set_xlabel("altitude")
        axs[0, 0].set_ylabel("photon density")
        axs[0, 0].legend()

        axs[1, 0].plot(
            alt,
            ds[i],
            "--",
            label=f"AWR: {degs[i]}" + "\N{DEGREE SIGN}",
            color=colors[i],
        )
        axs[1, 0].plot(
            alt,
            jds[i],
            "-",
            label=f"JFK: {degs[i]}" + "\N{DEGREE SIGN}",
            color=colors[i],
        )
        axs[1, 0].set_title("Log Photon Density")
        axs[1, 0].set_xlabel("altitude")
        axs[1, 0].set_ylabel("log photon density")
        axs[1, 0].set_yscale("log")
        axs[1, 0].legend()

        axs[0, 1].plot(
            alt,
            cs[i],
            "--",
            label=f"AWR: {degs[i]}" + "\N{DEGREE SIGN}",
            color=colors[i],
        )
        axs[0, 1].plot(
            alt,
            jcs[i],
            "-",
            label=f"JFK: {degs[i]}" + "\N{DEGREE SIGN}",
            color=colors[i],
        )
        axs[0, 1].set_title("Cherenkov Angle")
        axs[0, 1].set_xlabel("altitude")
        axs[0, 1].set_ylabel("Cherenkov Angle")
        axs[0, 1].legend()

        axs[1, 1].plot(
            alt,
            cs[i],
            "--",
            label=f"AWR: {degs[i]}" + "\N{DEGREE SIGN}",
            color=colors[i],
        )
        axs[1, 1].plot(
            alt,
            jcs[i],
            "-",
            label=f"JFK: {degs[i]}" + "\N{DEGREE SIGN}",
            color=colors[i],
        )
        axs[1, 1].set_title("Log Cherenkov Angle")
        axs[1, 1].set_xlabel("altitude")
        axs[1, 1].set_ylabel("Log Cherenkov Angle")
        axs[1, 1].set_yscale("log")
        axs[1, 1].legend()

    plt.show()
