import sys
import numpy as np
import colour  # Note: This is called colour-science in pip!
from scipy.interpolate import interp1d
from scipy.integrate import simpson
from scipy.optimize import curve_fit
import matplotlib.pyplot as plt
from cycler import cycler

# Fit complex refractive index per wavelength to per RGB color channel
# Usage: python fit_f82.py input
# where input is expected to be from https://refractiveindex.info

# We're dealing with reflectances, so convert to equal-energy illuminant
CIE_E = colour.CCS_ILLUMINANTS["CIE 1931 2 Degree Standard Observer"]["E"]

# Load NK data (expects three columns: Wavelength in µm, N, K), filter for visible wavelengths
nk_data = np.array([v for v in np.loadtxt(sys.argv[1]) if v[0] >= 0.36 and v[0] <= 0.8])

# Ground truth conductive Fresnel
def conductor(c, n, k):
    k2 = k * k
    rs_num = n * n + k2 - 2 * n * c + c * c
    rs_den = n * n + k2 + 2 * n * c + c * c
    rs = rs_num / rs_den

    rp_num = (n * n + k2) * c * c - 2 * n * c + 1
    rp_den = (n * n + k2) * c * c + 2 * n * c + 1
    rp = rp_num / rp_den

    return 0.5 * (rs + rp)


# Evaluate reflectivity for various incidence angles
reflectivity = {}
for cosI in np.linspace(0, 1, 101):
    # Evaluate Fresnel in spectral domain
    refl = colour.SpectralDistribution(
        conductor(cosI, nk_data[:, 1], nk_data[:, 2]), nk_data[:, 0] * 1000
    )
    xyz = colour.sd_to_XYZ(refl, method="Integration") / 100
    rgb = colour.XYZ_to_sRGB(xyz, illuminant=CIE_E, apply_cctf_encoding=False)
    reflectivity[cosI] = rgb

# Plot computed ground truth RGB curve
default_cycler = cycler(color=['r', 'g', 'b'])
plt.rc('axes', prop_cycle=default_cycler)
plt.plot(reflectivity.keys(), reflectivity.values(), marker=".", linestyle="None")

# Fit model to computed curves
param_fit = np.array(
    [
        curve_fit(
            conductor,
            np.array(list(reflectivity.keys())),
            [v[i] for v in reflectivity.values()],
            bounds=(0, np.inf)
        )[0]
        for i in range(3)
    ]
).transpose()

np.set_printoptions(suppress=True, precision = 4)
print('n:', param_fit[0])
print('k:', param_fit[1])

# Show fit curve
plt.plot(
    reflectivity.keys(),
    np.array(
        [conductor(cosi, param_fit[0], param_fit[1]) for cosi in reflectivity.keys()]
    ),
)
plt.xlabel(r'$\cos(i)$')
plt.ylabel('reflectivity')
plt.show()