#!/usr/bin/env python
"""
This is a script to calculate the index of refraction of water as a function of
wavelength, temperature, and pressure and fit it to a linear approximation in a
wavelength region where the PMTs are sensitive to. The data comes from [1].

[1] "Refractive index of Water and Steam as Function of Wavelength, Temperature, and Density". Schiebener et al. 1989.
"""
from __future__ import print_function, division
import numpy as np

A0 = 0.243905091
A1 = 9.53518094e-3
A2 = -3.64358110e-3
A3 = 2.65666426e-4
A4 = 1.59189325e-3
A5 = 2.45733798e-3
A6 = 0.897478251
A7 = -1.63066183e-2
UV = 0.2292020
IR = 5.432937

def get_index(p, wavelength, T):
    """
    Returns the index of pure water for a given density, wavelength, and
    temperature. The density should be in units of kg/m^3, the wavelength in nm,
    and the temperature in Celsius.
    """
    # normalize the density, temperature, and pressure
    p = p/1000.0
    wavelength = wavelength/589.0
    T = (T+273.15)/273.15

    # first we compute the right hand side of Equation 7
    c = A0 + A1*p + A2*T + A3*wavelength**2*T + A4/wavelength**2 + A5/(wavelength**2-UV**2) + A6/(wavelength**2-IR**2) + A7*p**2
    c *= p

    return np.sqrt((2*c+1)/(1-c))

if __name__ == '__main__':
    import matplotlib.pyplot as plt
    import argparse

    parser = argparse.ArgumentParser("fit a Taylor series approximation to the inverse of the index of refraction of water")
    parser.add_argument("--temperature", type=float, default=10.0)
    parser.add_argument("--density", type=float, default=1000.0)
    parser.add_argument("--deg", type=int, default=2)
    parser.add_argument("--x0", type=float, default=400.0)
    parser.add_argument("--low", type=float, default=300.0)
    parser.add_argument("--high", type=float, default=600.0)
    args = parser.parse_args()

    x = np.linspace(args.low,args.high,1000)
    n = get_index(args.density,x,args.temperature)
    z = np.polyfit(x-args.x0,1/n, args.deg)
    p = np.poly1d(z)

    print(p)

    plt.plot(x,1/n,label="data")
    plt.plot(x,p(x-args.x0),label="fit")
    plt.xlabel("Wavelength (nm)")
    plt.ylabel("1/n")
    plt.title("Index of Refraction of Water at 273.15 K")
    plt.legend()
    plt.show()