'''
---------------
Before running
---------------

This is the code for making the "Masking image" of FITS file.
The "Masking image" masks the 1) nearby stars, 2) Cosmic ray, (if this is MSI data) 3) the polarization mask area.

1.
Input file:  '*.fits'                      Preprocessed FITS file
             '*.mag.1'                     IRAF Phot file containing target's center info.
                                           See below (i.e.,2. What you need to run this code)
             'mask_*.fits'                 Masking image produced by 'Masking_image.py'.

Output file: 'result_Photo_*.csv'          Photometric result of each images
             'result_Pol_*.csv'            Polarimetric result of each sets

2.What you need to run this code.
  (The following packages must be installed.)
  i.  astropy (https://www.astropy.org/)
  iii.``*.mag.1`` file from IRAF's Phot package that contains the center of the target (2005 UD).
      The first line should be the center in the ordinary component, and the
      second line should be for extraordinary component.
      We've also uploaded the ``*.mag.1`` file that we used with this code.

3.
In this code, the center of the target is found by using the phot task of IRAF.
So, we need the ``.mag`` file to import the coordinate of target's ceter.
There is no problem if you find the target's center by other methods.
All you need to do is modifying the part that brings the central coordinate of target.
See ``BRING THE CENTER COORDINATE OF 2005 UD`` part.

4.
Directory should contain the complete sets consist of 4 images (taken at HWP=0, 22.5, 45, 67.5 deg).
If even one set does not have 4 images (e.g., set having images taken at HWP = 0, 45, 67.5 deg),
an error will occur.
'''

import glob
import os
import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
# from mpl_toolkits.axes_grid1 import make_axes_locatable
import pandas as pd

from astropy.io import fits, ascii
from astropy.time import Time
from astroquery.jplhorizons import Horizons
from astropy.modeling.models import Gaussian1D
from astropy.modeling.fitting import LevMarLSQFitter
from photutils import CircularAperture, CircularAnnulus, aperture_photometry
from astropy.stats import gaussian_fwhm_to_sigma

import warnings
from astropy.io.fits.verify import VerifyWarning
warnings.simplefilter('ignore', category=VerifyWarning)
mpl.rc('figure', max_open_warning=0)

# ********************************************************************************************************** #
# *                                           INPUT VALUE                                                  * #
# ********************************************************************************************************** #

Observatory = 'q33'  # Observatory code of Pirks
OBJECT = 155140  # 2005 UD

path = os.path.join('Directory/Where/image/to/be/analyzed/is/saved/')  # Path of directory where images to be analyzed are saved
IMAGE_list = glob.glob(os.path.join(path, '*.fits'))  # Bring all FITS preprocessed image in the directory

####################################
# Photometry Parameter
####################################
Aperture_scale = 1.7   # Aperture radius = Aperture_scale * FWHM
ANN_sacle = 4         # Annulus radius = ANN_sacle * FWHM
Dan = 20              # [pix] #Dannulus size

####################################
# Calibration Parameter (Please see Ishiguro et al. 2021)
####################################
eff = 0.9913  # Polarimetric efficiency of instrument
eff_err = 0.0001  # Error of polarimetric efficiency of instrument

q_inst = 0.00791  # Instrumental polarization, q_{inst}
u_inst = 0.00339  # Instrumental polarization, u_{inst}
eq_inst = 0.00025  # Error of instrumental polarization, error of q_{inst}
eu_inst = 0.00020  # Error of instrumental polarization, error of u_{inst}

theta_offset = 3.66  # Offset of polarization angle
err_theta = 0.17  # Error of offset

# ********************************************************************************************************** #

Photo_Log = pd.DataFrame({})
Pol_Log = pd.DataFrame({})
order = np.arange(0, len(IMAGE_list), 4)
IMAGE_list = sorted(IMAGE_list)

for z in order:
    print('==== Set {0:03d} ===='.format(int(z/4+1)))
    SET = [IMAGE_list[z], IMAGE_list[z+1], IMAGE_list[z+2], IMAGE_list[z+3]]  # Image taken at HWP=0,22.5,45,67.5

    I = []  # I_ex/I_o
    S = []  # (Sigma_{I_ex}/I_ex)**2 + (Sigma_{I_o}/I_o)**2
    INR_STR = []
    INR_END = []
    PSANG = []
    PA = []
    JD_mean = []
    file_name = SET[0].split('/')[-1].split('.')[0] + ' ~ '+SET[3].split('/')[-1][-7:]
    aper_seri_o = ''
    aper_seri_e = ''
    snr = 0

    # fig,ax = plt.subplots(4,2,figsize=(20,40))
    for ang in range(0, 4):
        image_i = SET[ang]
        hdul = fits.open(image_i)
        header = hdul[0].header
        data = hdul[0].data
        RET_ANG2 = header['RET-ANG2']

        # Bring the header info.==================================================
        iNR_STR = header['INR-STR']
        iNR_END = header['INR-END']
        INST_pa = header['INST-PA']
        JD = np.mean([header['MJD-STR'], header['MJD-END']])
        JD = JD + 2400000.5  # MJD -> JD
        epoch = Time(JD, format='jd').isot
        RN = header['RDNOISE']  # [electron]
        gain = header['GAIN']  # [electron/DN]
        Filter = header['FILTER']
        exp = header['EXPTIME']  # in sec

        # ************************************************************************************* #
        # *                Bring the observer quantities from JPL Horizons                    * #
        # ************************************************************************************* #
        obj = Horizons(id=OBJECT, location=Observatory, epochs=JD)
        eph = obj.ephemerides()
        psANG = eph['sunTargetPA'][0]  # [deg]
        pA = eph['alpha'][0]  # [deg]

        # ************************************************************************************* #
        # *                   BRING THE CENTER COORDINATE OF 2005 UD                          * #
        # ************************************************************************************* #
        # Here, we use Phot package in IRAF to find the center of the target
        # If you find the center of the target by other methods, please change this part.

        # Bring the center
        mag1 = ascii.read(image_i+'.mag.1')
        xo, yo = (mag1['XCENTER'][0]-1, mag1['YCENTER'][0]-1)  # pixel coordinate (x,y) of target's center in ordinary component
        xe, ye = (mag1['XCENTER'][1]-1, mag1['YCENTER'][1]-1)  # pixel coordinate (x,y) of target's center in extra. component

        # ************************************************************************************* #
        # *                            Bring the Masking Image                                * #
        # ************************************************************************************* #
        masking_path = os.path.join(path, 'Masking', 'mask_'+image_i.split('/')[-1])
        hdul_mask = fits.open(masking_path)[0]
        masking = hdul_mask.data

        # ************************************************************************************* #
        # *                            Determine FWHM of target                               * #
        # ************************************************************************************* #
        def crop(data, row, col, size):
            row_str = int(row-size)
            row_end = int(row+size)
            col_str = int(col-size)
            col_end = int(col+size)
            data_cr = data[row_str:row_end, col_str:col_end]
            return data_cr
        ##############################
        #  Ordinary
        ##############################
        mdata = np.ma.masked_array(data, masking)
        crop_image = crop(mdata, yo, xo, 35)
        sum_crop_image = np.mean(crop_image, axis=0)
        sum_crop_image = sum_crop_image - np.mean(sum_crop_image[:10])

        g_init = Gaussian1D(amplitude=np.mean(sum_crop_image[34:36]),
                            mean=35,
                            stddev=10*gaussian_fwhm_to_sigma,
                            bounds={'mean': (33, 37),
                                    'stddev': ((5*gaussian_fwhm_to_sigma, 20*gaussian_fwhm_to_sigma))})

        x = np.arange(0, len(crop_image), 1)
        fitter = LevMarLSQFitter()
        fitted = fitter(g_init, x, sum_crop_image)

        re_g_init = Gaussian1D(amplitude=fitted.amplitude.value,
                               mean=fitted.mean.value,
                               stddev=fitted.stddev.value,
                               bounds={'mean': (33, 37),
                                       'stddev': ((5*gaussian_fwhm_to_sigma, 20*gaussian_fwhm_to_sigma))})
        fitter = LevMarLSQFitter()
        fitted = fitter(re_g_init, x, sum_crop_image)
        FWHM_ordi = 2*fitted.stddev.value*np.sqrt(2*np.log(2))
        aper_seri_o += '{0:.1f} '.format(FWHM_ordi)

        ##############################
        #  Extra-Ordinary
        ##############################
        crop_image = crop(mdata, ye, xe, 35)
        sum_crop_image = np.mean(crop_image, axis=0)
        sum_crop_image = sum_crop_image - np.mean(sum_crop_image[:10])
        g_init = Gaussian1D(amplitude=np.mean(sum_crop_image[34:36]),
                            mean=35,
                            stddev=10*gaussian_fwhm_to_sigma,
                            bounds={'mean': (33, 37),
                                    'stddev': (5*gaussian_fwhm_to_sigma, 20*gaussian_fwhm_to_sigma)})

        fitter = LevMarLSQFitter()
        fitted = fitter(g_init, x, sum_crop_image)

        re_g_init = Gaussian1D(amplitude=fitted.amplitude.value,
                               mean=fitted.mean.value,
                               stddev=fitted.stddev.value,
                               bounds={'mean': (33, 37),
                                       'stddev': (5*gaussian_fwhm_to_sigma, 20*gaussian_fwhm_to_sigma)})
        fitter = LevMarLSQFitter()
        fitted = fitter(re_g_init, x, sum_crop_image)
        FWHM_extra = 2*fitted.stddev.value*np.sqrt(2*np.log(2))
        aper_seri_e += '{0:.1f} '.format(FWHM_extra)

        # ************************************************************************************* #
        # *                        Circular Aperture Photometry                               * #
        # ************************************************************************************* #
        def skyvalue(data, y0, x0, r_in, r_out, masking):
            masking = masking.astype(bool)
            ann = CircularAnnulus([x0, y0], r_in=r_in, r_out=r_out)
            phot = aperture_photometry(data, ann, mask=masking)

            phot1 = aperture_photometry(masking*1, ann)
            pixel_count = phot1['aperture_sum'][0]
            sky = phot['aperture_sum'][0]/(ann.area-pixel_count)

            masked_image = np.ma.masked_array(data, masking)

            # Determine sky std
            y_in = int(y0-r_out)
            y_out = int(y0+r_out)
            x_in = int(x0-r_out)
            x_out = int(x0+r_out)
            if y_in < 0:
                y_in = 0
            if y_out < 0:
                y_out = 0
            if x_in < 0:
                x_in = 0
            if x_out < 0:
                x_out = 0
            masked_image = masked_image[y_in:y_out, x_in:x_out]
            new_mask = np.zeros(np.shape(masked_image))+1
            for yi in range(len(masked_image)):
                for xi in range(len(masked_image[0])):
                    position = (xi - r_out)**2 + (yi-r_out)**2
                    if position < (r_out)**2 and position > r_in**2:
                        new_mask[yi, xi] = 0
            new_mask = new_mask.astype(bool)
            Sky_region = np.ma.masked_array(masked_image, new_mask)
            std = np.ma.std(Sky_region)
            return(sky, std, ann.area-pixel_count)

        def signal_to_noise(source_eps, sky_std, rd, npix):
            signal = source_eps
            noise = np.sqrt( (source_eps  + npix *
                              (sky_std**2 )) + npix * rd ** 2)
            return signal / noise

        FWHM_sel = max(FWHM_ordi, FWHM_extra)  # Selected FWHM
        Aperture_radius = Aperture_scale*FWHM_sel/2
        Annulus_radius = ANN_sacle*FWHM_sel/2
        Ann_out = Annulus_radius+Dan

        # Determine sky value by aperture
        Aper_o = CircularAperture([xo, yo], Aperture_radius)  # Set aperture
        sky_o, sky_std_o, area_o = skyvalue(data, yo, xo, Annulus_radius, Ann_out, masking)  # Set area determinung Sk #[count]

        Aper_e = CircularAperture([xe, ye], Aperture_radius)  # Set aperture
        sky_e, sky_std_e, area_e = skyvalue(data, ye, xe, Annulus_radius, Ann_out, masking)  # Set area determinung Sk

        sky_std_o = sky_std_o*gain
        sky_std_e = sky_std_e*gain

        ########
        masking = masking.astype(bool)
        masked_image = np.ma.masked_array(data, masking)
        y0, x0 = yo, xo
        y_in = int(y0-Ann_out)
        y_out = int(y0-Annulus_radius)
        x_in = int(x0-Ann_out)
        x_out = int(x0-Annulus_radius)
        crop_masked_images = masked_image[y_in:y_out, x_in:x_out]
        std = np.ma.std(crop_masked_images)

        # Target aperture sum
        Flux_o = aperture_photometry(data - sky_o, Aper_o, masking)['aperture_sum'][0]*gain  # in e-
        ERR_o = np.sqrt(Flux_o + 3.14*Aperture_radius**2*(sky_std_o**2 + (RN*gain)**2))  # in e-
        sky_o = sky_o*gain  # in e-
        Snr_o = signal_to_noise(Flux_o, sky_std_o, RN, Aperture_radius**2*3.14)

        Flux_e = aperture_photometry(data - sky_e, Aper_e, masking)['aperture_sum'][0]*gain  # in e-
        ERR_e = np.sqrt(Flux_e + 3.14*Aperture_radius**2*(sky_std_e**2 + (RN*gain)**2))  # in e-
        sky_e = sky_e*gain  # in e-
        Snr_e = signal_to_noise(Flux_e, sky_std_e, RN, Aperture_radius**2*3.14)
        snr += Snr_o + Snr_e

        # ************************************************************************************* #
        # *              Record the values for calculating Stokes parameter                   * #
        # ************************************************************************************* #
        I.append(Flux_e/Flux_o)
        S_ret = (ERR_e/(Flux_e))**2 + (ERR_o/(Flux_o))**2
        S.append(S_ret)
        PSANG.append(psANG)
        PA.append(pA)
        JD_mean.append(JD)
        INR_STR.append(iNR_STR)
        INR_END.append(iNR_END)

        # ************************************************************************************* #
        # *                       Plot the aperture photometry                                * #
        # ************************************************************************************* #
        def circle(x, y, r):
            theta = np.linspace(0, 2*np.pi, 100)
            x1 = r*np.cos(theta)+y
            x2 = r*np.sin(theta)+x
            return(x2.tolist(), x1.tolist())

        plot_data = np.ma.masked_array(data, masking)
        figsize = 50
        # im = ax[ang,0].imshow(plot_data - sky_o/gain,vmin=-50,vmax=50,cmap='seismic')
        # xi,yi = circle(xo,yo,Aperture_radius)
        # ax[ang,0].plot(xi,yi,color='y',lw=2)
        # xi,yi = circle(xo,yo,Annulus_radius)
        # ax[ang,0].plot(xi,yi ,color='c',lw=2)
        # xi,yi = circle(xo,yo,Annulus_radius+Dan)
        # ax[ang,0].plot(xi,yi ,color='c',lw=2)
        # ax[ang,0].plot(xo,yo,marker='X',color='c',ms=4)
        # ax[ang,0].set_xlim(xo-figsize,xo+figsize)
        # ax[ang,0].set_ylim(yo-figsize,yo+figsize)
        # ax[ang,0].set_title('Ordinary'+image_i.split('/')[-1],fontsize=14)
        # divider = make_axes_locatable(ax[ang,0])
        # cax = divider.append_axes("right", size="5%", pad=0.05)
        # plt.colorbar(im,cax=cax)

        # im = ax[ang,1].imshow(plot_data - sky_e/gain,vmin=-50,vmax=50,cmap='seismic')
        # xi,yi = circle(xe, ye,Aperture_radius)
        # ax[ang,1].plot(xi,yi,color='y',lw=2)
        # xi,yi = circle(xe, ye,Annulus_radius)
        # ax[ang,1].plot(xi,yi ,color='c',lw=2)
        # xi,yi = circle(xe, ye,Annulus_radius+Dan)
        # ax[ang,1].plot(xi,yi ,color='c',lw=2)
        # ax[ang,1].plot(xe,ye,marker='X',color='c',ms=4)
        # ax[ang,1].set_xlim(xe-figsize,xe+figsize)
        # ax[ang,1].set_ylim(ye-figsize,ye+figsize)
        # ax[ang,1].set_title('ExtaOrdinary'+image_i.split('/')[-1],fontsize=14)
        # divider = make_axes_locatable(ax[ang,1])
        # cax = divider.append_axes("right", size="5%", pad=0.05)
        # plt.colorbar(im,cax=cax)

        # ************************************************************************************* #
        # *                          Save result of photometry                                * #
        # ************************************************************************************* #
        print(image_i.split('/')[-1])
        Photo_Log = Photo_Log.append({'Object': header['OBJECT'],
                                      'Filename': image_i.split('/')[-1],
                                      'set': int(z/4)+1,
                                      'HWPANG': header['RET-ANG2'],
                                      'TIME': epoch,
                                      'JD': np.mean(JD_mean),
                                      'FWHM_o': aper_seri_o,
                                      'FWHM_e': aper_seri_e,
                                      'Aper [pix]': Aperture_radius,
                                      'EXP [s]': exp,
                                      'Ann': Annulus_radius,
                                      'Ann_out': Ann_out,
                                      'Flux_o': Flux_o,
                                      'eFLux_o': ERR_o,
                                      'Flux_e': Flux_e,
                                      'eFLux_e': ERR_e,
                                      'SNR_o': Snr_o,
                                      'SNR_e': Snr_e,
                                      'Sky_o': sky_o,
                                      'eSky_o': sky_std_o,
                                      'Sky_e': sky_e,
                                      'eSky_e': sky_std_e}, ignore_index=True)

    # ************************************************************************************* #
    # *                         Calculate Stokes Parameter                                * #
    # ************************************************************************************* #
    PsANG_av = np.mean(PSANG)  # the average position angles of Sun-target radius vector of a set
    PA_av = np.mean(PA)  # the average phase angle of a set
    JD_av = np.mean(JD_mean)  # the average JD of a set

    I0 = I[0]
    I45 = I[1]
    I22 = I[2]
    I67 = I[3]

    S0 = S[0]
    S45 = S[1]
    S22 = S[2]
    S67 = S[3]

    Rq = np.sqrt(I0/I45)
    Ru = np.sqrt(I22/I67)

    q = (Rq - 1)/(Rq + 1)
    u = (Ru - 1)/(Ru + 1)

    q_ran = Rq/((Rq + 1)**2)  * np.sqrt(S0 + S45)
    u_ran = Ru/((Ru + 1)**2)  * np.sqrt(S22 + S67)

    q_sys = 0
    u_sys = 0

    q_err = np.sqrt(q_ran**2 + q_sys**2)
    u_err = np.sqrt(u_ran**2 + u_sys**2)

    # ====================
    # Correct Efficiency
    # ====================
    qq = q/eff
    uu = u/eff

    # random error of corrected q,u
    qq_ran = q_ran/eff
    uu_ran = u_ran/eff

    # the systematic errors
    qq_sys = np.abs(q)*eff_err/eff
    uu_sys = np.abs(u)*eff_err/eff

    # ====================
    # Correc Instrumental polarization
    # ====================

    STR0, STR45, STR22, STR67 = INR_STR[0], INR_STR[1], INR_STR[2], INR_STR[3]
    END0, END45, END22, END67 = INR_END[0], INR_END[1], INR_END[2], INR_END[3]

    # averaged value of frame-reader value
    rq = np.deg2rad((STR0 + END0 + STR45 + END45)/4.)   # Instrument star, end value ( 0,0,45,45 ) in rad
    ru = np.deg2rad((STR22 + END22 + STR67 + END67)/4.)  # Instrument star, end value ( 22.5,22.5,67.5,67.5) in rad

    qqq = qq - ((q_inst * np.cos(2*rq)) - (u_inst * np.sin(2*rq)))
    uuu = uu - ((q_inst * np.sin(2*ru)) + (u_inst * np.cos(2*ru)))

    # random error of corrected q,u
    qqq_ran = qq_ran
    uuu_ran = uu_ran

    # the systematic errors
    qqq_sys = np.sqrt( qq_sys**2 + (eq_inst * np.cos(2*rq))**2 +
                       (eu_inst*np.sin(2*rq))**2 )
    uuu_sys = np.sqrt( uu_sys**2 + (eq_inst * np.sin(2*ru))**2 +
                       (eu_inst*np.cos(2*ru))**2 )

    # ====================
    # Transform_CelestialCoord
    # ====================
    theta = np.deg2rad(theta_offset)
    the_err = np.deg2rad(err_theta)

    qqqq = qqq * np.cos(2*theta) + uuu*np.sin(2*theta)
    uuuu = -qqq * np.sin(2*theta) + uuu*np.cos(2*theta)

    qqqq_ran = np.sqrt( (qqq_ran*np.cos(2*theta))**2 + (uuu_ran*np.sin(2*theta))**2 )
    uuuu_ran = np.sqrt( (qqq_ran*np.sin(2*theta))**2 + (uuu_ran*np.cos(2*theta))**2 )

    qqqq_sys = np.sqrt( (qqq_sys*np.cos(2*theta))**2 +
                        (uuu_sys*np.sin(2*theta))**2 +
                        (np.pi/180*2*uuuu*the_err)**2 )
    uuuu_sys = np.sqrt( (qqq_sys*np.sin(2*theta))**2 +
                        (uuu_sys*np.cos(2*theta))**2 +
                        (np.pi/180*2*qqqq*the_err)**2 )

    q, q_ran, q_sys, u, u_ran, u_sys = qqqq, qqqq_ran, qqqq_sys, uuuu, uuuu_ran, uuuu_sys

    # ************************************************************************************* #
    # *                       Calculate Polarimetric result                               * #
    # ************************************************************************************* #
    P = np.sqrt(q**2 + u**2)
    P_ran = np.sqrt( (q*q_ran)**2 + (u*u_ran)**2 )/P
    P_sys = np.sqrt( (q*q_sys)**2 + (u*u_sys)**2 )/P
    theta_pol = np.rad2deg(1/2 * np.arctan2(u, q))

    # Random noise correction (Wardle & Kronberg 1974)
    if P**2 >= P_ran**2:
        print('Random error bias correction is done.')
        P_cor = np.sqrt(P**2 - P_ran**2)
    elif P**2 < P_ran**2 :
        print('Due to P < randome error, random error bias correction is NOT done.')
        P_cor = 0

    P_error = np.sqrt(P_ran**2 + P_sys**2)  # Polarization error

    if P_cor != 0:
        ran_PolAng = 1/2 * 180/3.14 * P_ran/P_cor
        sys_PolAng = 1/2 * 180/3.14 * P_sys/P_cor
        PolAng_error = np.sqrt(ran_PolAng**2 + sys_PolAng**2)
    elif P_cor == 0:
        ran_PolAng = 51.96
        sys_PolAng = 51.96
        PolAng_error = 51.96
        # Naghizadeh-Khouei & Clarke 1993

    if PsANG_av + 90 < 180:
        pi = PsANG_av + 90
    else:
        pi = PsANG_av - 90

    # Converted a polarization degree with respect to the scattering plane
    theta_r = theta_pol - pi
    Pr = P_cor * np.cos(2*np.deg2rad(theta_r))
    Pol_Log = Pol_Log.append({'filename': file_name,
                              'Filter': 'Rc',
                              'JD': JD_av,
                              'alpha [deg]': PA_av,
                              'PsANG [deg]': PsANG_av,
                              'FWHM_o': aper_seri_o,
                              'FWHM_e': aper_seri_e,
                              'Aper_radius [pix]': Aperture_radius,
                              'q': q,
                              'u': u,
                              'ran_q': q_ran,
                              'ran_u': u_ran,
                              'sys_q': q_sys,
                              'sys_u': u_sys,
                              'PsANG': PsANG_av,
                              'theta': theta_pol,
                              'theta_r': theta_r,
                              'eTheta': PolAng_error,
                              'P': P_cor,
                              'eP': np.sqrt(P_ran**2 + P_sys**2)},
                             ignore_index=True)
    plt.show()

# ************************************************************************************* #
# *                       Calculate the weighted mean                                 * #
# ************************************************************************************* #

def weight(x, err):
    x = np.array(x)
    err = np.array(err)

    w = 1/err**2
    sumW = np.sum(w)
    weight = w/sumW

    xav = np.sum(weight*x)
    Err = 1/np.sqrt(sumW)

    return(xav, Err)

time_ = Time(Pol_Log['JD'][0], format='jd').iso
q_av, ranq_av = weight(Pol_Log['q'], Pol_Log['ran_q'])
u_av, ranu_av = weight(Pol_Log['u'], Pol_Log['ran_u'])
sysq_av = np.mean(Pol_Log['sys_q'])
sysu_av = np.mean(Pol_Log['sys_u'])
errq_av = (ranq_av**2 + sysq_av**2)**0.5
erru_av = (ranu_av**2 + sysu_av**2)**0.5

P = np.sqrt(q_av**2+u_av**2)
ran_P = np.sqrt((q_av*ranq_av)**2 + (u_av*ranu_av)**2)/P
sys_P = np.sqrt((q_av*sysq_av)**2 + (u_av*sysu_av)**2)/P
eP = np.sqrt(ran_P**2 + sys_P**2)

# Random noise correction (Wardle & Kronberg 1974)
if P**2 - ran_P**2 < 0:
    print('Random error bias correction is done.')
    Pcor = 0
else:
    print('Due to P < random error, P = 0% ')
    Pcor = np.sqrt(P**2 - ran_P**2)

# Converted a polarization degree with respect to the scattering plane
theta = 1/2*np.rad2deg(np.arctan2(u_av, q_av))
psang = np.mean(Pol_Log['PsANG'])
if psang+90 > 180:
    pi = psang-90
else:
    pi = psang + 90

if Pcor != 0:
    theta_ran = 1/2*180/3.14*ran_P/Pcor
    theta_sys = 1/2*180/3.14*sys_P/Pcor
    eTheta = np.sqrt(theta_ran**2 + theta_sys**2)
elif Pcor == 0:
    theta_ran = 51.96
    theta_sys = 51.96
    eTheta = 51.96
    # Naghizadeh-Khouei & Clarke 1993

Pol_Log = Pol_Log.append({'filename': 'Weighted_average',
                          'JD': np.mean(Pol_Log['JD'].values),
                          'alpha [deg]': np.mean(Pol_Log['alpha [deg]'].values),
                          'PsANG [deg]': np.mean(Pol_Log['PsANG [deg]'].values),
                          'q': q_av,
                          'u': u_av,
                          'ran_q': ranq_av,
                          'ran_u': ranu_av,
                          'sys_q': sysq_av,
                          'sys_u': sysu_av,
                          'theta': theta,
                          'theta_r': theta_r,
                          'P': Pcor,
                          'eP': eP,
                          'Aper_radius [pix]': np.mean(Pol_Log['Aper_radius [pix]'].values),
                          'Pr': P_cor*np.cos(np.deg2rad(theta_r*2)),
                          'eTheta': eTheta
                          },
                         ignore_index=True)

Pol_name = ['filename', 'JD', 'alpha [deg]', 'PsANG [deg]',
            'q', 'u', 'ran_q', 'ran_u', 'sys_q', 'sys_u',
            'P', 'eP', 'Pr', 'theta', 'theta_r', 'eTheta',
            'Aper_radius [pix]']
Phot_name = ['Object', 'Filename', 'set', 'TIME', 'JD',
             'HWPANG', 'FWHM_o', 'FWHM_e', 'EXP [s]',
             'Aper [pix]', 'Ann', 'Ann_out', 'Flux_o',
             'eFLux_o', 'Flux_e', 'eFLux_e', 'Sky_o',
             'eSky_o', 'Sky_e', 'eSky_e', 'SNR_o',
             'SNR_e']

Pol_Log = Pol_Log.reindex(columns=Pol_name)
Photo_Log = Photo_Log.reindex(columns=Phot_name)
Pol_Log = Pol_Log.round({'JD': 6, 'alpha [deg]': 2, 'PsANG [deg]': 2,
                         'q': 4, 'u': 4, 'ran_q': 4, 'ran_u': 4,
                         'sys_q': 4, 'sys_u': 4, 'theta': 2, 'theta_r': 2,
                         'P': 4, 'eP': 4, 'Aper_radius [pix]': 2, 'Pr': 4,
                         'eTheta': 2})
Photo_Log = Photo_Log.round({'HWPANG': 2, 'JD': 6, 'FWHM_o': 2, 'FWHM_e': 2,
                             'Aper [pix]': 2, 'EXP [s]': 2, 'Ann': 2,
                             'Ann_out': 2, 'Flux_o': 4, 'eFLux_o': 4,
                             'Flux_e': 4, 'eFLux_e': 4, 'SNR_o': 2,
                             'SNR_e': 2, 'Sky_o': 4, 'eSky_o': 4,
                             'Sky_e': 4, 'eSky_e': 4})
Pol_Log.to_csv(os.path.join(path, 'result_Pol_{}.csv'.format(time_.split(' ')[0])))
Photo_Log.to_csv(os.path.join(path, 'result_Photo_{}.csv'.format(time_.split(' ')[0])))
print(os.path.join(path, 'result_Pol_{}.csv'.format(time_.split(' ')[0])) + ' is created.')
print(os.path.join(path, 'result_Photo_{}.csv'.format(time_.split(' ')[0])) + ' is created.')


# MIT License
#
# Copyright (c) 2017 Ulrich Noebauer
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
"""
Module providing a number of basic calculators to determine the steady-state
structure of line-driven hot star winds analytically. In particular, the
predictions according to
 * Castor, Abbott and Klein 1975
 * Friend and Abbott 1986
 * Kudritzki, Pauldrach, Puls and Abbott 1989
are included. Throughout this module, we rely heavily on the standard text
book on stellar winds by Lamers and Cassinelli 1999.

Note
----
We assume a so-called frozen-in ionization throughout the wind in all the wind
structure calculators (i.e. delta=0 in terms of the CAK force multipliers).

References
----------
 * Castor, Abbott and Klein 1975 (CAK75)
   'Radiation-driven winds in Of stars', Astrophysical Journal,
   1975, 195, 157-174
 * Friend and Abbott 1986 (FA86)
   'The theory of radiatively driven stellar winds. III - Wind models with
   finite disk correction and rotation', Astrophysical Journal,
   1986, 311, 701-707
 * Kudritzki, Pauldrach, Puls and Abbott 1989 (KPPA89)
   'Radiation-driven winds of hot stars. VI - Analytical solutions for wind
   models including the finite cone angle effect', Astronomy & Astrophysics,
   1989, 219, 205-218
 * Noebauer and Sim 2015 (NS15)
   'Self-consistent modelling of line-driven hot-star winds with Monte Carlo
   radiation hydrodynamics', Monthly Notices of the Royal Astronomical Society,
   2015, 453, 3120-3134
 * Lamers and Cassinelli 1999 (LC99)
   'Introduction to Stellar Winds', Cambridge University Press, 1999
"""
from __future__ import print_function
import numpy as np
import scipy.integrate as integ
import astropy.units as units
import astropy.constants as csts
from astropy.utils.decorators import lazyproperty

def _test_unit(val, final_unit_string, initial_unit_string):
    """helper routine to add unit to a quantity and convert to different units

    If val is not yet an astropy.units.quantity, it is assumed that it is
    given in units specified by the initial_unit_string. The input is then
    returned after converting to the units given by the final_unit_string.

    Parameters
    ----------
    val : float, int, np.ndarray
        scalar or vector input, can either already be an astropy quantity or
        not
    final_unit_string : str
        string describing the desired final units for the input
    initial_unit_string : str
        string describing the assumed initial units of the input

    Returns
    -------
    res : astropy.units.quantity
        input converted to units given by final_unit_string
    """
    try:
        val.to(final_unit_string)
    except AttributeError:
        val = val * units.Unit(initial_unit_string)
    res = val.to(final_unit_string)

    return res

class StarBase(object):
    """Base class containing the fundamental properties of the central star

    Parameters
    ---------
    mass : float, astropy.units.quantity
        stellar mass, if dimensionless, it is assumed to be in units of solar
        masses (default 52.5)
    lum : float, astropy.units.quantity
        stellar luminosity, if dimensionless, it is assumed to be in units of
        solar luminosities (default 1e6)
    teff : float, astropy.units.quantity
        effective temperature, if dimensionless, it is assumed to be Kelvin
        (default 4.2e4)
    gamma : float
        Eddington factor due to electron scattering (default 0)
    sigma : float, astropy.units.quantity
        reference specific electron scattering cross section, if dimensionless
        it is assumed to be in units of cm^2/g (default 0.3)
    """
    def __init__(self, mass=52.5, lum=1e6, teff=4.2e4, gamma=0, sigma=0.3):

        self.mass = mass
        self.lum = lum
        self.teff = teff
        self.sigma = sigma
        self.gamma = gamma

    @property
    def mass(self):
        return self._mass

    @mass.setter
    def mass(self, val):
        val = _test_unit(val, "g", "solMass")
        self._mass = val

    @property
    def lum(self):
        return self._lum

    @lum.setter
    def lum(self, val):
        val = _test_unit(val, "erg/s", "solLum")
        self._lum = val

    @property
    def teff(self):
        return self._teff

    @teff.setter
    def teff(self, val):
        val = _test_unit(val, "K", "K")
        self._teff = val

    @property
    def sigma(self):
        return self._sigma

    @sigma.setter
    def sigma(self, val):
        val = _test_unit(val, "cm^2/g", "cm^2/g")
        self._sigma = val

    @lazyproperty
    def rad(self):
        """stellar radius"""
        rad = np.sqrt(self.lum /
                      (4 * np.pi * csts.sigma_sb * self.teff**4))
        return rad

    @lazyproperty
    def vth(self):
        """thermal velocity (see LC99, eqs. 8.8, 8.83)"""
        vth = np.sqrt(2. * self.teff * csts.k_B / csts.u)
        return vth

    @lazyproperty
    def vesc(self):
        """escape velocity from stellar surface, accounting for electron
        scattering (see LC99, eq. 2.39)"""
        vesc = np.sqrt(2. * csts.G * self.mass * (1. - self.gamma) / self.rad)
        return vesc

class WindBase(object):
    """Base class containing the fundamental properties of the wind

    Parameters
    ----------
    alpha : float
        CAK force multiplier parameter, see LC99, eq. 8.86 (default 0.6)
    k : float
        CAK force multiplier parameter, see LC99, eq. 8.86 (default 0.5)
    """
    def __init__(self, alpha=0.6, k=0.5):

        self.alpha = alpha
        self.k = k

class WindStructureBase(object):
    """Base class describing the basic structure of a star+wind system

    Parameters
    ----------
    mstar : float, astropy.units.quantity
        stellar mass, see StarBase for details (default 52.5)
    lstar : float, astropy.units.quantity
        stellar luminosity, see StarBase for details (default 1e6)
    teff : float, astropy.units.quantity
        effective temperature, see StarBase for details (default 4.2e4)
    alpha : float
        force multiplier parameter, see WindBase for details (default, 0.6)
    k : float
        force multiplier parameter, see WindBase for details (default, 0.5)
    gamma : float
        Eddington factor, see StarBase for details (default 0)
    sigma : float, astropy.units.quantity
        reference electron scattering cross section, see StarBase for details
        (default 0.3)
    """
    def __init__(self, mstar=52.5, lstar=1e6, teff=4.2e4, alpha=0.6, k=0.5,
                 gamma=0, sigma=0.3):

        self.star = StarBase(mass=mstar, lum=lstar, teff=teff, gamma=gamma,
                             sigma=sigma)
        self.wind = WindBase(alpha=alpha, k=k)

    @lazyproperty
    def eta(self):
        """wind efficiency (see LC99, eq. 8.20)"""
        eta = (self.mdot * self.vterm * csts.c / self.star.lum).to("").value
        return eta

    @lazyproperty
    def t(self):
        """CAK dimensionless optical depth, c.f. LC99, eq. 8.82 and 8.104"""
        t = (self.mdot * self.star.sigma * self.star.vth /
             (2. * np.pi * self.vterm**2 * self.star.rad)).to("")
        return t

    @lazyproperty
    def m(self):
        """CAK force multiplier, assuming frozen-in ionization, c.f. LC99, eq.
        8.86 with delta=0"""
        m = self.wind.k * self.t**(-self.wind.alpha)
        return m

class BaseVelocityDensityMixin(object):
    """Mixin class providing routines to calculate the wind velocity and
    density structure"""
    def v(self, x):
        """calculate wind velocity according to CAK at given location

        C.f. LC99, eq. 8.104

        Parameters
        ----------
        x : float, np.ndarray
            dimensionless position, i.e. r/Rstar

        Returns
        -------
        v : float, np.ndarry
            wind velocity
        """

        return (self.vterm * np.sqrt(1. - 1. / x)).to("km/s")

    def rho(self, x):
        """calculate wind density at given location

        C.f. LC99, eq. 3.1

        Parameters
        ----------
        x : float, np.ndarray
            dimensionless position, i.e. r/Rstar

        Returns
        -------
        rho : float, np.ndarry
            wind density
        """

        r = self.star.rad * x

        return (self.mdot /
                (4. * np.pi * r**2 * self.v(x))).to("g/cm^3")

class BaseCakStructureMixin(object):
    """Mixin class providing the CAK mass loss rate and terminal wind speed"""

    @lazyproperty
    def mdot_cak(self):
        """Mass-loss rate according to CAK75, see LC99, eq. 8.105"""

        mdot_cak = ((4. * np.pi / (self.star.sigma * self.star.vth)) *
                    (self.star.sigma / (4. * np.pi))**(1. / self.wind.alpha) *
                    ((1. - self.wind.alpha) / self.wind.alpha)**(
                        (1. - self.wind.alpha) / self.wind.alpha) *
                    (self.wind.alpha * self.wind.k)**(1. / self.wind.alpha) *
                    (self.star.lum / csts.c)**(1. / self.wind.alpha) *
                    (csts.G * self.star.mass * (1. - self.star.gamma))**(
                        (self.wind.alpha - 1.) / self.wind.alpha))
        return mdot_cak.to("Msun/yr")

    @lazyproperty
    def vterm_cak(self):
        """Terminal wind speed according to CAK75, see LC99, eq. 8.104"""

        vterm_cak = (np.sqrt(self.wind.alpha / (1. - self.wind.alpha)) *
                     self.star.vesc)

        return vterm_cak.to("km/s")

class WindStructureCak75(WindStructureBase, BaseCakStructureMixin,
                         BaseVelocityDensityMixin):
    """Wind Structure Calculator based on the approach by CAK75.

    The wind structure is determined based on a CAK line-driving force,
    assuming a frozen-in ionization state and a central point source.

    Paramters
    ---------
    see WindStructureBase
    """
    def __init__(self, mstar=52.5, lstar=1e6, teff=4.2e4, alpha=0.6, k=0.5,
                 gamma=0, sigma=0.3):
        super(WindStructureCak75, self).__init__(mstar=mstar, lstar=lstar,
                                                 teff=teff, alpha=alpha,
                                                 k=k, gamma=gamma,
                                                 sigma=sigma)

    @property
    def mdot(self):
        """mass loss rate, equal to basic CAK mass loss rate"""
        return self.mdot_cak.to("Msun/yr")

    @property
    def vterm(self):
        """terminal wind speed, equal to basic CAK terminal velocity"""
        return self.vterm_cak.to("km/s")

class WindStructureKppa89(WindStructureBase, BaseCakStructureMixin,
                          BaseVelocityDensityMixin):
    """Wind Structure Calculator based on the approach by KPPA89.

    The wind structure is determined based on a CAK line-driving force,
    assuming a frozen-in ionization state but taking the finite size of the
    central star into account.

    Paramters
    ---------
    see WindStructureBase
    beta : float
        exponent for the beta-type velocity law (default 0.8, see LC99, sec.
        8.9.2 iii)
    """
    def __init__(self, mstar=52.5, lstar=1e6, teff=4.2e4, alpha=0.6, k=0.5,
                 gamma=0, sigma=0.3, beta=0.8):
        super(WindStructureKppa89, self).__init__(mstar=mstar, lstar=lstar,
                                                  teff=teff, alpha=alpha,
                                                  k=k, gamma=gamma,
                                                  sigma=sigma)

        self.f1 = 1. / (self.wind.alpha + 1.)
        self.beta = beta

    @lazyproperty
    def mdot(self):
        """mass loss rate, see KPPA89, eq. 31"""

        mdot = self.f1**(1. / self.wind.alpha) * self.mdot_cak

        return mdot.to("Msun/yr")

    @lazyproperty
    def vterm(self):
        """terminal wind speed, see KPPA89, eq. 39"""

        vterm = self.vterm_cak * np.sqrt(integ.quad(self.z, 0, 1)[0])

        return vterm.to("km/s")

    def h(self, x):
        """see KPPA89, eq. 14"""

        return (x - 1.) / self.beta

    def f(self, x):
        """see KPPA89, eqs. 16, 15"""

        return (1. / (self.wind.alpha + 1.) * x**2 / (1. - self.h(x)) *
                (1. - (1. - 1. / x**2 + self.h(x) / x**2)**(
                    self.wind.alpha + 1.)))

    def fn(self, x):
        """see KPPA89, eq. 35"""

        return self.f(x) / self.f1

    def z(self, u):
        """see KPPA89, eq. 36"""

        x = 1. / u
        z = (self.fn(x)**(1. / (1. - self.wind.alpha)) *
             (1. + np.sqrt(2. / self.wind.alpha *
                           (1. - (1. / self.fn(x))**(
                               1. / (1. - self.wind.alpha))))))
        return z

    def _v_scalar(self, x):
        """see KPPA89, eq. 36"""

        u = 1. / x
        I = integ.quad(self.z, u, 1)[0]
        vesc2 = (2. * csts.G * self.star.mass *
                 (1. - self.star.gamma) / self.star.rad)
        v = np.sqrt(self.wind.alpha / (1. - self.wind.alpha) * vesc2 * I).to(
            "km/s")

        return v.value

    def v(self, x):
        """calculate wind velocity according to KPPA89 at given location

        Parameters
        ----------
        x : float, np.ndarray
            dimensionless position, i.e. r/Rstar

        Returns
        -------
        v : float, np.ndarry
            wind velocity
        """

        if type(x) is np.ndarray:
            v = np.array([self._v_scalar(xi) for xi in x])
        else:
            v = self._v_scalar(x)

        v = v * units.km / units.s
        return v.to("km/s")

class WindStructureFa86(WindStructureBase, BaseCakStructureMixin,
                        BaseVelocityDensityMixin):
    """Wind Structure Calculator based on the approach by FA86.

    The wind structure is determined based on a CAK line-driving force,
    assuming a frozen-in ionization state but taking the finite size of the
    central star into account. All expressions for the wind properties result
    from fits to the numerical simulations as presented by FA86.

    Paramters
    ---------
    see WindStructureBase
    """
    def __init__(self, mstar=52.5, lstar=1e6, teff=4.2e4, alpha=0.6, k=0.5,
                 gamma=0, sigma=0.3):
        super(WindStructureFa86, self).__init__(mstar=mstar, lstar=lstar,
                                                teff=teff, alpha=alpha,
                                                k=k, gamma=gamma,
                                                sigma=sigma)

    @lazyproperty
    def mdot(self):
        """see FA86, eq. 9"""

        mdot = (self.mdot_cak * 0.5 *
                (self.star.vesc / (1e3 * units.km / units.s))**(-0.3))

        return mdot.to("Msun/yr")

    @lazyproperty
    def vterm(self):
        """see FA86, eq. 8"""

        vterm = (self.star.vesc * 2.2 * self.wind.alpha /
                 (1. - self.wind.alpha) *
                 (self.star.vesc / (1e3 * units.km / units.s))**0.2)

        return vterm.to("km/s")

    def v(self, x):
        """calculate wind velocity according to KPPA89 at given location

        See FA86, eq. 11

        Parameters
        ----------
        x : float, np.ndarray
            dimensionless position, i.e. r/Rstar

        Returns
        -------
        v : float, np.ndarry
            wind velocity
        """
        v = self.vterm * (1. - 1. / x)**(0.8)

        return v.to("km/s")

def example():
    """Example application of the wind structure calculators.

    The parameters are adopted from NS15, table 3 and are appropriate for the
    O star zeta-Puppis.
    """

    import matplotlib
    import os
    if "DISPLAY" not in os.environ:
        matplotlib.use("TkAgg")
    import matplotlib.pyplot as plt
    plt.rcParams["text.usetex"] = False

    mstar = 52.5 * units.solMass
    lstar = 1e6 * units.solLum
    teff = 4.2e4 * units.K
    sigma = 0.3 * units.cm**2 / units.g

    k = 0.381
    alpha = 0.595
    gamma = 0.502

    x = np.logspace(-2, 2, 512) + 1

    wind_cak75 = WindStructureCak75(mstar=mstar, lstar=lstar, teff=teff, k=k,
                                    alpha=alpha, gamma=gamma, sigma=sigma)
    wind_fa86 = WindStructureFa86(mstar=mstar, lstar=lstar, teff=teff, k=k,
                                  alpha=alpha, gamma=gamma, sigma=sigma)
    wind_kppa89 = WindStructureKppa89(mstar=mstar, lstar=lstar, teff=teff, k=k,
                                      alpha=alpha, gamma=gamma, sigma=sigma)

    fig = plt.figure(figsize=(10, 10))
    fig.subplots_adjust(wspace=0.3, hspace=0.3)

    plt.subplot(221)
    plt.plot(x, wind_cak75.v(x), label="CAK75")
    plt.plot(x, wind_fa86.v(x), ls="dashed", label="FA86")
    plt.plot(x, wind_kppa89.v(x), ls="dashdot", label="KPPA86")
    plt.yscale("log")
    plt.ylim([1e2, 3e3])
    plt.xlim([0.8, 10])
    plt.xlabel(r"$r/R_{\star}$")
    plt.ylabel(r"$v$ [km/s]")
    plt.legend(frameon=False)

    plt.subplot(222)
    plt.plot(x - 1, wind_cak75.v(x) / wind_cak75.vterm, label="CAK75")
    plt.plot(x - 1, wind_fa86.v(x) / wind_fa86.vterm, ls="dashed",
             label="FA86")
    plt.plot(x - 1, wind_kppa89.v(x) / wind_kppa89.vterm, ls="dashdot",
             label="KPPA86")
    plt.xscale("log")
    plt.xlim([1e-2, 1e2])
    plt.ylim([0, 1])
    plt.xlabel(r"$r/R_{\star} - 1$")
    plt.ylabel(r"$v/v_{\infty}$")

    plt.subplot(223)
    plt.plot(x, wind_cak75.rho(x), label="CAK75")
    plt.plot(x, wind_fa86.rho(x), ls="dashed", label="FA86")
    plt.plot(x, wind_kppa89.rho(x), ls="dashdot", label="KPPA86")
    plt.yscale("log")
    plt.ylim([1e-15, 1e-10])
    plt.xlim([0.8, 10])
    plt.xlabel(r"$r/R_{\star}$")
    plt.ylabel(r"$\rho$ $[\mathrm{g\,cm^{-3}}]$")

    plt.subplot(224)
    plt.plot(x, wind_cak75.mdot * np.ones(len(x)))
    plt.plot(x, wind_fa86.mdot * np.ones(len(x)), ls="dashed")
    plt.plot(x, wind_kppa89.mdot * np.ones(len(x)), ls="dashdot")
    plt.yscale("log")
    plt.xlim([0.8, 10])
    plt.ylim([1e-5, 1e-4])
    plt.xlabel(r"$r/R_{\star}$")
    plt.ylabel(r"$\dot M$ $[\mathrm{M_{\odot}\,yr^{-1}}]$")
    fig.savefig("wind_structure_example.pdf")
    plt.show()

if __name__ == "__main__":

    example()

import numpy,import scipy,import astropy,from astropy


# Python 3.5
# Script for MetalSilicate python library
# Copyright Kayla Iacovino
# This script contains pandas dataframes with activity coefficient data digitized from the Steelmaking Data Source Book

import pandas as pd

# These hard-coded pandas DataFrames are generated from an Excel Spreadsheet and hardcoded for easier access.
"""
The code to generate these:
import pandas as pd
import numpy as np
activity_coeffs_df = pd.read_excel('ActivityCoeffs_and_InteractionParams.xlsx', sheet_name='ActivityCoefficients')
activity_coeffs_df.fillna("No Data")

import math
f = open("hard_coded_activity_coefficients.py","w+")
cols = ['i', 'Element_State', 'Fe_State', 'gamma_naught_i', 'Temp_K', 'DeltaG_J_per_g-atom', 'TempRange_K', 'Ref', 'Year', 'Note']

f.write ("\n")
f.write("activity_coeffs = pd.DataFrame({\n")
for col in cols:
	iterno = 1
	f.write("'" + (str(col)+"': ["))
	for index, row in activity_coeffs_df.iterrows():
		try:
			value = float(row[col])
			if str(value) == 'nan':
				f.write("'No Data'")
			else:
				f.write(str(value))
		except:
			f.write("'" + str(row[col]) + "'")
		if iterno < len(activity_coeffs_df.index):
			f.write(",")
		iterno += 1
	f.write("], \n")
f.write(" }).set_index('i') \n")

##NOTE! Delete the final comma at the end of the final column.

##Reminder: to look up a value, use syntax: activity_coeffs.loc['Ag']['gamma_naught_i']
"""

activity_coeffs = pd.DataFrame({
'i': ['Ag','Al','Al_delta','B','C','C_austenite','C_ferrite','Ca','Ce','Co','Cr','Cr_solid','Cu','H','H_delta','H_austenite','H_ferrite','La','Mn','Mo','Mo_solid','N','N_delta','N_austenite','N_ferrite','Nb','Nb','Ni','O','P','Pb','S','Si','Sn','Ta','Ti','Ti_solid','U','V','V_solid','W','W_solid','Zr','Zr_solid'],
'Element_State': ['l','l','l','s','gr','gr','gr','l','l','l','l','s','l','g','g','g','g','l','l','l','s','g','g','g','g','l','s','l','g','g','l','g','l','l','l','l','s','l','l','s','l','s','l','s'],
'Fe_State': ['liquid iron','liquid iron','delta iron','liquid iron','liquid iron','austenite','ferrite','liquid iron','liquid iron','liquid iron','liquid iron','liquid iron','liquid iron','liquid iron','delta iron','austenite','ferrite','liquid iron','liquid iron','liquid iron','liquid iron','liquid iron','delta iron','austenite','ferrite','liquid iron','liquid iron','liquid iron','liquid iron','liquid iron','liquid iron','liquid iron','liquid iron','liquid iron','liquid iron','liquid iron','liquid iron','liquid iron','liquid iron','liquid iron','liquid iron','liquid iron','liquid iron','liquid iron'],
'gamma_naught_i': [200.0,0.049,0.027,0.022,0.538,7.49,313.0,2270.0,0.322,0.55,1.0,1.14,8.58,'–','–','–','–',9.3,1.44,1.0,2.2,'–','–','–','–',0.2,1.4,0.66,'–','–',837.0,'–',0.0013,2.58,0.04,0.004,0.009,0.027,0.08,0.1,1.0,7.6,0.037,0.043],
'Temp_K': [1873.0,1873.0,1673.0,1873.0,1873.0,1273.0,1073.0,1880.0,1873.0,1873.0,1873.0,1873.0,1873.0,1873.0,1773.0,1273.0,1073.0,1873.0,1843.0,1823.0,1823.0,1873.0,1773.0,1273.0,1073.0,1873.0,1873.0,1873.0,1873.0,1873.0,1873.0,1873.0,1873.0,1873.0,1873.0,1873.0,1873.0,1873.0,1873.0,1873.0,1873.0,1873.0,1873.0,1873.0],
'DeltaG_J_per_g-atom': ['82400-43.76T','-17100-19.4T','-17900-19.3T','-65300-21.55T','17230-39.87T','72180+52.05T*logT-227.03T','83720-55.76T','No Data','-16700-46.44T','No Data','-37.70T','19200-46.86T','33500-39.37T','36460+30.46T','28650+45.35T','22630+45.35T','28650+45.35T','125,900-94.6T','No Data','No Data','No Data','9916+20.17T','29090+19.91T','-8620+37.42T','29090+19.1T','No Data','23000-52.3T','-20900-31.1T','-117100-3.39T','-157700+5.4T','No Data','-125100+18.5T','-131500-17.24T','No Data','No Data','No Data','No Data','-56100-50.3T','-42300-36.0T','-20700-45.6T','-48.1T','No Data','-51000-42.38T','-34700-50.00T'],
'TempRange_K': ['No Data','m.p.–1873','1673–m.p.','No Data','1773–2073','1773–1673','923–1073','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data','No Data'],
'Ref': [2.0,6.0,6.0,2.0,90.0,76.0,100.0,17.0,102.0,203.0,2.0,2.0,2.0,22.0,3434.0,434.0,434.0,282.0,288.0,436.0,436.0,437.0,438.0,439.0,438.0,440.0,2.0,2.0,360.0,441.0,41.0,442.0,2.0,513.0,440.0,440.0,385.0,2.0,2.0,2.0,2.0,436.0,2.0,2.0],
'Year': [1974.0,1977.0,1977.0,1974.0,1962.0,1970.0,1967.0,1964.0,1977.0,1978.0,1974.0,1974.0,1974.0,1963.0,1950.0,1950.0,1950.0,1977.0,1982.0,1984.0,1984.0,1982.0,1973.0,1955.0,1973.0,1973.0,1974.0,1974.0,1959.0,1980.0,1971.0,1981.0,1974.0,1981.0,1973.0,1973.0,1984.0,1974.0,1974.0,1974.0,1974.0,1984.0,1974.0,1974.0],
'Note': ['From solubility data, regular solution assumed','From solubility data, regular solution assumed','No Data','No Data','No Data','No Data','No Data','No Data','No Data','log(gammaCO) = -0.257N^2_Fe','Ideal solution assumed','No Data','No Data','1/2H2','1/2H2','1/2H2','1/2H2','No Data','No Data','No Data','±0.17','1/2N2','1/2N2','1/2N2','1/2N2','No Data','No Data','No Data','1/2O2; Data were combined with deltaG_0 = -251877+58.325T for reaction H2(g) + 1/2O2(g) = H2O(g) [14]','1/2P2','No Data','1/2S2','No Data','Mass analysis','Calculated value','No Data','Calculated value','Regular solution assumed','No Data','No Data','No Data','±1.1','Gamma_0 Zr assumed equal to Gamma_0 Ti and regular solution assumed','No Data']
 }).set_index('i')


'''
A class that schedules scans submitted by clients.

Author: Nathan Rowley
Date: August 2018
'''

from astropy.coordinates import SkyCoord, EarthLocation, AltAz
from astropy.time import Time
from astropy import units as u
from srtutility.NTPTime import NTPTime
import sqlite3
import re
import datetime
import pytz
from astral import Astral

class Schedule:

	# Initializes a Schedule instance for building scan schedules.
	#
	# :param starttime: start time of the schedule period in unix time
	# :param endtime: end time of the schedule period in unix time
	def __init__(self, starttime, endtime):

		startblock = self.Block(None, starttime, starttime)			# create Blocks to mark the start and end of the schedule

		endblock = self.Block(None, endtime, endtime)

		self.schedule = [startblock, endblock]

		self.localtz = pytz.timezone('America/Chicago')		# set local timezone for display time conversions

		self.database_location = '../srtdatabase/srtdata.db'

	# Internal class for storing scan and time info in a Schedule
	#
	# :param scanid: the id of the scan in this block, or None if the block is an endmarker
	# :param starttime: start time of the scan in unix time
	# :param endtime: end time of the scan in unix time
	class Block:

		def __init__(self, scanid, starttime, endtime):

			self.scanid = scanid
			self.starttime = starttime
			self.endtime = endtime

	# Method for adding a scan to the schedule. Inserts the scan at the earliest possible valid time.
	#
	# :param scanid: the id of the scan to be added to the schedule
	# :param curtime: the current unix time
	# :return status: the status of the scan indicating success or failure to schedule
	def schedulescan(self, scanid, curtime):

		srtdb = sqlite3.connect(self.database_location)			# establish a connection and cursor into the database
		srtdb.row_factory = sqlite3.Row
		cur = srtdb.cursor()

		scanparams = cur.execute("SELECT * FROM SCANPARAMS WHERE ID = ?", (scanid,)).fetchone()

		duration = re.split('[hms]', scanparams['duration'])								# get duration values of scan

		seconds = int(duration[0]) * 60 * 60 + int(duration[1]) * 60 + int(duration[2])		# calculate durationin seconds

		if curtime > self.schedule[0].endtime:

			starttime = curtime + 300						# start time is the current time if past the start of night

		else:

			starttime = self.schedule[0].endtime + 300		# start time of new Block is after the first Block

		endtime = starttime + seconds

		scantype = scanparams['type']

		print(str(starttime) + ' ' + str(endtime) + ' ' + scanparams['ras'] + ' ' + scanparams['dec'] + ' ' + scantype)

		status = 'durationerror'

		print('attempting to schedule a scan')

		for i in range(len(self.schedule) - 1):					# loops through the spaces between each Block in the schedule

			print('checking for space after block ' + str(i))

			while starttime >= self.schedule[i].endtime + 300 and endtime <= self.schedule[i+1].starttime - 300:	# loops through five-minute steps of the time between two blocks (padded by five minutes on either side)

				print('there is space, checking for validity')

				result = self.checkscan(scanparams['ras'], scanparams['dec'], scantype, starttime, endtime)		# check to see if the scan is valid within a particular time window

				print(result)

				if status == 'durationerror' or status == 'positionerror':		# error hierarchy is movebounds > position > duration

					status = result

				if result == 'scheduled':		# if the scan is valid, create a new Block for the scan and insert it into the schedule

					print('found a valid spot!')

					status = result

					newblock = Schedule.Block(scanid, starttime, endtime)

					starttime = datetime.datetime.fromtimestamp(starttime, pytz.utc).astimezone(self.localtz).strftime('%H:%M')	# convert time values to local time strings
					endtime = datetime.datetime.fromtimestamp(endtime, pytz.utc).astimezone(self.localtz).strftime('%H:%M')

					cur.execute("INSERT INTO SCHEDULE VALUES (?,?,?)", (scanid, starttime, endtime))	# update the schedule and scan status in the db
					srtdb.commit()

					srtdb.close()

					self.schedule.insert(i+1, newblock)

					return status

				starttime += 300			# if scan is not valid within the time frame, adjust the frame five minutes forward
				endtime += 300

			starttime = self.schedule[i+1].endtime + 300	# update starttime to after the next Block

			if starttime < curtime:							# if curtime not yet reached, reset starttime to curtime

				starttime = curtime + 300

			endtime = starttime + seconds					# update endtime to reflect new starttime

		print('couldn\'t find a spot. error code: ' + status)

		srtdb.close()

		return status

	# Method that removes a scan from the schedule with an id matching scanid.
	#
	# :param scanid: the id of the scan to be removed
	# :return:
	def deschedulescan(self, scanid):

		for i in range(len(self.schedule - 2)):

			if self.schedule[i].scanid == scanid:

				print('scan removed from schedule')

				del self.schedule[i]
				break

	# Helper method for checking the validity of a scan.
	# Checks that the scan's position is in the sky and that the scan does not go out of movement bounds.
	#
	# :param ras: right ascension of the scan to be checked
	# :param dec: declination of the scan to be checked
	# :param scantype: a string designating the type of scan
	# :param starttime: start time of the scan in unix time
	# :param endtime: end time of the scan in unix time
	# :return: a string indicating the status of the scan
	def checkscan(self, ras, dec, scantype, starttime, endtime):

		srtdb = sqlite3.connect(self.database_location)	# establish a connection and cursor into the database
		srtdb.row_factory = sqlite3.Row
		cur = srtdb.cursor()

		configdata = cur.execute("SELECT * FROM CONFIG").fetchone()		# retrieve config data from the database

		position = SkyCoord(ras, dec, frame = 'icrs')					# convert position into astropy SkyCoord object for coord transformation

		location = EarthLocation(lat = configdata['lat'], lon = configdata['lon'], height = configdata['height']) 	# convert location into astropy EarthLocation

		srtdb.close()

		middletime = (starttime + endtime) / 2

		starttime = Time(starttime, format = 'unix')						# convert times into astropy Time objects
		startframe = AltAz(location = location, obstime = starttime)		# create astropy AltAz frame objects for each time

		if scantype == 'track':

			middletime = Time(middletime, format = 'unix')						# necessary to check a midpoint for altitude bounds due to circumpolar positions possibly dipping too low
			middleframe = AltAz(location = location, obstime = middletime)
			endtime = Time(endtime, format = 'unix')
			endframe = AltAz(location = location, obstime = endtime)

		try:														# attempt to transform the position to the AltAz frame of each time

			startpos = position.transform_to(startframe)

			if scantype == 'track':

				middlepos = position.transform_to(middleframe)
				endpos = position.transform_to(endframe)

		except ValueError as e:										# if the position can't be transformed, return a position error

			return 'positionerror'

		azbounds = (configdata['azlower'], configdata['azupper'])		# repackage movement bounds
		albounds = (configdata['allower'], configdata['alupper'])

		startaz = float(startpos.az.to_string(unit=u.deg, decimal=True))
		startal = float(startpos.alt.to_string(unit=u.deg, decimal=True))

		if scantype == 'track':

			middleaz = float(middlepos.az.to_string(unit=u.deg, decimal=True))
			middleal = float(middlepos.alt.to_string(unit=u.deg, decimal=True))
			endaz = float(endpos.az.to_string(unit=u.deg, decimal=True))
			endal = float(endpos.alt.to_string(unit=u.deg, decimal=True))

		if startal < 0 or startal > 180:		# if a position has a negative altitude, object is not in the sky, so return position error

			return 'positionerror'

		if scantype == 'track' and (middleal < 0 or middleal > 180 or endal < 0 or endal > 180):

			return 'positionerror'

		valid = True 													# check that the scan stays within telscope movement bounds

		valid = valid and startaz >= azbounds[0] and startaz <= azbounds[1]
		valid = valid and startal >= albounds[0] and startal <= albounds[1]

		if scantype == 'track':

			valid = valid and middleaz >= azbounds[0] and middleaz <= azbounds[1]
			valid = valid and middleal >= albounds[0] and middleal <= albounds[1]
			valid = valid and endaz >= azbounds[0] and endaz <= azbounds[1]
			valid = valid and endal >= albounds[0] and endal <= albounds[1]

		if not valid:

			print(str(startaz) + ' ' + str(startal))

			return 'moveboundserror'

		return 'scheduled'

def main():

	srtdb = sqlite3.connect('../srtdatabase/srtdata.db')		# establish a connection and cursor into the database
	srtdb.row_factory = sqlite3.Row
	cur = srtdb.cursor()

	config = cur.execute("SELECT * FROM CONFIG").fetchone()

	a = Astral()

	today = datetime.date.today()

	nighttimes = a.night_utc(today, config['lat'], config['lon'])

	print(nighttimes[0].isoformat() + ' ' + nighttimes[1].isoformat())

	dusk = nighttimes[0].timestamp()
	dawn = nighttimes[1].timestamp()

	schedule = Schedule(dusk, dawn)

	cur.execute("DELETE FROM SCHEDULE")

	# cur.execute("INSERT INTO SOURCES VALUES (?,?,?)", ('sigma octantis', '21h8m47s', '-88d57m23s'))
	# cur.execute("INSERT INTO SCANIDS VALUES (?,?,?)", (-30, 'invalidtest', 'submitted'))
	# cur.execute("INSERT INTO SCANPARAMS VALUES (?,?,?,?,?,?,?,?,?)", (-30, 'track', 'sigma octantis', '21h8m47s', '-88d57m23s', '1h0m0s', 1500, 1510, 10))
	# cur.execute("UPDATE CONFIG SET LAT = ?, LON = ?", (44.45, -93.16))
	# cur.execute("INSERT INTO SOURCES VALUES (?,?,?)", ('polaris', '2h31m49s', '89d15m50s'))
	# cur.execute("INSERT INTO SCANIDS VALUES (?,?,?)", (-20, 'scheduletest', 'submitted'))
	# cur.execute("INSERT INTO SCANPARAMS VALUES (?,?,?,?,?,?,?,?,?)", (-20, 'track', 'polaris', '2h31m49s', '89d15m50s', '0h30m0s', 1500, 1510, 10))
	# srtdb.commit()

	ntp = NTPTime()

	schedule.schedulescan(-20, ntp.getcurrenttime())
	schedule.schedulescan(-30, ntp.getcurrenttime())

	for block in schedule.schedule:

		print(str(block.scanid) + ' ' + str(block.starttime) + ' ' + str(block.endtime))

# main()


# -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2012-2016 GEM Foundation
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake. If not, see <http://www.gnu.org/licenses/>.

"""
Module exports :class:`AtkinsonBoore2006`,
:class:`AtkinsonBoore2006MblgAB1987bar140NSHMP2008`,
:class:`AtkinsonBoore2006MblgJ1996bar140NSHMP2008`,
:class:`AtkinsonBoore2006Mwbar140NSHMP2008`,
:class:`AtkinsonBoore2006MblgAB1987bar200NSHMP2008`,
:class:`AtkinsonBoore2006MblgJ1996bar200NSHMP2008`,
:class:`AtkinsonBoore2006Mwbar200NSHMP2008`,
:class:`AtkinsonBoore2006Modified2011`.
"""
from __future__ import division

import numpy as np
# standard acceleration of gravity in m/s**2
from scipy.constants import g
from math import log10

from openquake.hazardlib.gsim.boore_atkinson_2008 import BooreAtkinson2008
from openquake.hazardlib.gsim.utils import (
    mblg_to_mw_atkinson_boore_87,
    mblg_to_mw_johnston_96,
    clip_mean
)
from openquake.hazardlib.gsim.base import CoeffsTable
from openquake.hazardlib import const
from openquake.hazardlib.imt import PGA, PGV, SA

class AtkinsonBoore2006(BooreAtkinson2008):
    """
    Implements GMPE developed by Gail M. Atkinson and David M. Boore and
    published as "Earthquake Ground-Motion Prediction Equations for Eastern
    North America" (2006, Bulletin of the Seismological Society of America,
    Volume 96, No. 6, pages 2181-2205). This class implements only the
    equations for stress parameter of 140 bars. The correction described in
    'Adjustment of Equations to Consider Alternative Stress Parameters',
    p. 2198, is not implemented.
    This class extends the BooreAtkinson2008 because it uses the same soil
    amplification function. Note that in the paper, the reported soil
    amplification function is the one used in a preliminary version of the
    Boore and Atkinson 2008 GMPE, while the one that should be used is the
    one described in the final paper. See comment in:
    http://www.daveboore.com/pubs_online/ab06_gmpes_programs_and_tables.pdf
    """
    #: Supported tectonic region type is stable continental, given
    #: that the equations have been derived for Eastern North America
    DEFINED_FOR_TECTONIC_REGION_TYPE = const.TRT.STABLE_CONTINENTAL

    #: Supported intensity measure types are spectral acceleration,
    #: peak ground velocity and peak ground acceleration, see paragraph
    #: 'Methodology and Model Parameters', p. 2182
    DEFINED_FOR_INTENSITY_MEASURE_TYPES = set([
        PGA,
        PGV,
        SA
    ])

    #: Supported intensity measure component is horizontal
    #: :attr:`~openquake.hazardlib.const.IMC.HORIZONTAL`,
    #: see paragraph 'Results', pag 2190, and caption to table 6, p. 2192
    DEFINED_FOR_INTENSITY_MEASURE_COMPONENT = const.IMC.HORIZONTAL

    #: Supported standard deviation type is total, see table 6
    #: and 9, p. 2192 and 2202, respectively.
    DEFINED_FOR_STANDARD_DEVIATION_TYPES = set([
        const.StdDev.TOTAL
    ])

    #: Required site parameters is Vs30.
    #: See paragraph 'Equations for soil sites', p. 2200
    REQUIRES_SITES_PARAMETERS = set(('vs30', ))

    #: Required rupture parameter is magnitude (see
    #: paragraph 'Methodology and Model Parameters', p. 2182)
    REQUIRES_RUPTURE_PARAMETERS = set(('mag', ))

    #: Required distance measure is Rrup.
    #: See paragraph 'Methodology and Model Parameters', p. 2182
    REQUIRES_DISTANCES = set(('rrup', ))

    def get_mean_and_stddevs(self, sites, rup, dists, imt, stddev_types):
        """
        See :meth:`superclass method
        <.base.GroundShakingIntensityModel.get_mean_and_stddevs>`
        for spec of input and result values.
        """
        mean = self._get_mean(sites.vs30, rup.mag, dists.rrup, imt, scale_fac=0)
        stddevs = self._get_stddevs(stddev_types, num_sites=sites.vs30.size)

        return mean, stddevs

    def _get_mean(self, vs30, mag, rrup, imt, scale_fac):
        """
        Compute and return mean
        """
        C_HR, C_BC, C_SR, SC = self._extract_coeffs(imt)

        rrup = self._clip_distances(rrup)

        f0 = self._compute_f0_factor(rrup)
        f1 = self._compute_f1_factor(rrup)
        f2 = self._compute_f2_factor(rrup)

        pga_bc = self._get_pga_bc(
            f0, f1, f2, SC, mag, rrup, vs30, scale_fac
        )

        # compute mean values for hard-rock sites (vs30 >= 2000),
        # and non-hard-rock sites (vs30 < 2000) and add soil amplification
        # term
        mean = np.zeros_like(vs30)
        self._compute_mean(C_HR, f0, f1, f2, SC, mag, rrup,
                           vs30 >= 2000.0, mean, scale_fac)
        self._compute_mean(C_BC, f0, f1, f2, SC, mag, rrup,
                           vs30 < 2000.0, mean, scale_fac)
        self._compute_soil_amplification(C_SR, vs30, pga_bc, mean)

        # convert from base 10 to base e
        if imt == PGV():
            mean = np.log(10 ** mean)
        else:
            # convert from cm/s**2 to g
            mean = np.log((10 ** mean) * 1e-2 / g)

        return mean

    def _get_pga_bc(self, f0, f1, f2, SC, mag, rrup, vs30, scale_fac):
        """
        Compute and return PGA on BC boundary
        """
        pga_bc = np.zeros_like(vs30)
        self._compute_mean(self.COEFFS_BC[PGA()], f0, f1, f2, SC, mag,
                           rrup, vs30 < 2000.0, pga_bc, scale_fac)

        return (10 ** pga_bc) * 1e-2 / g

    def _extract_coeffs(self, imt):
        """
        Extract dictionaries of coefficients specific to required
        intensity measure type.
        """
        C_HR = self.COEFFS_HARD_ROCK[imt]
        C_BC = self.COEFFS_BC[imt]
        C_SR = self.COEFFS_SOIL_RESPONSE[imt]
        SC = self.STRESS_COEFFS[imt]

        return C_HR, C_BC, C_SR, SC

    def _clip_distances(self, rrup):
        """
        Return array of distances with values clipped to 1. See end of
        paragraph 'Methodology and Model Parameters', p. 2182. The equations
        have a singularity for distance = 0, so that's why distances are
        clipped to 1.
        """
        rrup = rrup.copy()
        rrup[rrup < 1] = 1

        return rrup

    def _compute_f0_factor(self, rrup):
        """
        Compute and return factor f0 - see equation (5), 6th term, p. 2191.
        """
        # f0 = max(log10(R0/rrup),0)
        f0 = np.log10(self.COEFFS_IMT_INDEPENDENT['R0'] / rrup)
        f0[f0 < 0] = 0.0

        return f0

    def _compute_f1_factor(self, rrup):
        """
        Compute and return factor f1 - see equation (5), 4th term, p. 2191
        """
        # f1 = min(log10(rrup),log10(R1))
        f1 = np.log10(rrup)
        logR1 = np.log10(self.COEFFS_IMT_INDEPENDENT['R1'])
        f1[f1 > logR1] = logR1

        return f1

    def _compute_f2_factor(self, rrup):
        """
        Compute and return factor f2, see equation (5), 5th term, pag 2191
        """
        # f2 = max(log10(rrup/R2),0)
        f2 = np.log10(rrup / self.COEFFS_IMT_INDEPENDENT['R2'])
        f2[f2 < 0] = 0.0

        return f2

    def _compute_stress_drop_adjustment(self, SC, mag, scale_fac):
        """
        Compute equation (6) p. 2200
        """
        return scale_fac * np.minimum(
            SC['delta'] + 0.05,
            0.05 + SC['delta'] * (
                np.maximum(mag - SC['M1'], 0) / (SC['Mh'] - SC['M1'])
            )
        )

    def _compute_mean(self, C, f0, f1, f2, SC, mag, rrup, idxs, mean,
                      scale_fac):
        """
        Compute mean value (for a set of indexes) without site amplification
        terms. This is equation (5), p. 2191, without S term.
        """
        mean[idxs] = (C['c1'] +
                      C['c2'] * mag +
                      C['c3'] * (mag ** 2) +
                      (C['c4'] + C['c5'] * mag) * f1[idxs] +
                      (C['c6'] + C['c7'] * mag) * f2[idxs] +
                      (C['c8'] + C['c9'] * mag) * f0[idxs] +
                      C['c10'] * rrup[idxs] +
                      self._compute_stress_drop_adjustment(SC, mag, scale_fac))

    def _compute_soil_amplification(self, C, vs30, pga_bc, mean):
        """
        Compute soil amplification, that is S term in equation (5), p. 2191,
        and add to mean values for non hard rock sites.
        """
        # convert from base e (as defined in BA2008) to base 10 (as used in
        # AB2006)
        sal = np.log10(np.exp(self._get_site_amplification_linear(vs30, C)))
        sanl = np.log10(np.exp(
            self._get_site_amplification_non_linear(vs30, pga_bc, C)))

        idxs = vs30 < 2000.0
        mean[idxs] = mean[idxs] + sal[idxs] + sanl[idxs]

    def _get_stddevs(self, stddev_types, num_sites):
        """
        Return total standard deviation (see table 6, p. 2192).
        """
        assert all(stddev_type in self.DEFINED_FOR_STANDARD_DEVIATION_TYPES
                   for stddev_type in stddev_types)
        stddevs = [np.zeros(num_sites) +
                   self.COEFFS_IMT_INDEPENDENT['std_total']
                   for _ in stddev_types]
        return stddevs

    #: Hard rock coefficents, table 6, pag 2192,
    #: coefficient values taken from Fortran implementation of Dave Boore
    #: (higher precision than in the paper)
    COEFFS_HARD_ROCK = CoeffsTable(sa_damping=5, table="""\
    IMT     c1          c2          c3          c4          c5          c6          c7          c8          c9          c10
    5.000  -5.408E+00   1.714E+00  -9.012E-02  -2.537E+00   2.267E-01  -1.268E+00   1.162E-01   9.792E-01  -1.767E-01  -1.757E-04
    4.000  -5.791E+00   1.916E+00  -1.071E-01  -2.441E+00   2.113E-01  -1.162E+00   1.018E-01   1.012E+00  -1.824E-01  -2.010E-04
    3.125  -6.038E+00   2.080E+00  -1.221E-01  -2.367E+00   2.002E-01  -1.073E+00   8.950E-02   1.002E+00  -1.803E-01  -2.306E-04
    2.500  -6.169E+00   2.211E+00  -1.348E-01  -2.299E+00   1.898E-01  -9.860E-01   7.860E-02   9.683E-01  -1.765E-01  -2.823E-04
    2.000  -6.183E+00   2.302E+00  -1.442E-01  -2.223E+00   1.770E-01  -9.370E-01   7.067E-02   9.518E-01  -1.768E-01  -3.220E-04
    1.587  -6.043E+00   2.342E+00  -1.496E-01  -2.157E+00   1.662E-01  -8.704E-01   6.047E-02   9.207E-01  -1.734E-01  -3.748E-04
    1.250  -5.724E+00   2.324E+00  -1.505E-01  -2.104E+00   1.565E-01  -8.202E-01   5.186E-02   8.563E-01  -1.661E-01  -4.329E-04
    1.000  -5.272E+00   2.264E+00  -1.483E-01  -2.069E+00   1.497E-01  -8.132E-01   4.666E-02   8.262E-01  -1.622E-01  -4.862E-04
    0.794  -4.604E+00   2.132E+00  -1.406E-01  -2.062E+00   1.468E-01  -7.974E-01   4.345E-02   7.748E-01  -1.558E-01  -5.790E-04
    0.629  -3.917E+00   1.987E+00  -1.314E-01  -2.045E+00   1.419E-01  -7.818E-01   4.297E-02   7.878E-01  -1.590E-01  -6.948E-04
    0.500  -3.216E+00   1.826E+00  -1.201E-01  -2.018E+00   1.344E-01  -8.134E-01   4.437E-02   8.839E-01  -1.751E-01  -7.704E-04
    0.397  -2.437E+00   1.649E+00  -1.084E-01  -2.051E+00   1.363E-01  -8.426E-01   4.483E-02   7.386E-01  -1.557E-01  -8.509E-04
    0.315  -1.721E+00   1.483E+00  -9.739E-02  -2.080E+00   1.382E-01  -8.893E-01   4.869E-02   6.101E-01  -1.389E-01  -9.538E-04
    0.251  -1.121E+00   1.342E+00  -8.722E-02  -2.082E+00   1.349E-01  -9.714E-01   5.628E-02   6.140E-01  -1.432E-01  -1.055E-03
    0.199  -6.153E-01   1.227E+00  -7.886E-02  -2.087E+00   1.312E-01  -1.120E+00   6.788E-02   6.055E-01  -1.459E-01  -1.125E-03
    0.158  -1.455E-01   1.123E+00  -7.143E-02  -2.116E+00   1.302E-01  -1.303E+00   8.311E-02   5.617E-01  -1.438E-01  -1.182E-03
    0.125   2.144E-01   1.054E+00  -6.664E-02  -2.154E+00   1.295E-01  -1.608E+00   1.046E-01   4.273E-01  -1.303E-01  -1.153E-03
    0.100   4.797E-01   1.017E+00  -6.404E-02  -2.201E+00   1.270E-01  -2.007E+00   1.326E-01   3.371E-01  -1.266E-01  -1.047E-03
    0.079   6.906E-01   9.974E-01  -6.276E-02  -2.262E+00   1.246E-01  -2.487E+00   1.636E-01   2.139E-01  -1.207E-01  -8.469E-04
    0.063   9.109E-01   9.802E-01  -6.208E-02  -2.360E+00   1.263E-01  -2.972E+00   1.910E-01   1.069E-01  -1.173E-01  -5.786E-04
    0.050   1.105E+00   9.719E-01  -6.197E-02  -2.466E+00   1.276E-01  -3.390E+00   2.144E-01  -1.391E-01  -9.839E-02  -3.167E-04
    0.040   1.264E+00   9.680E-01  -6.232E-02  -2.581E+00   1.317E-01  -3.644E+00   2.276E-01  -3.506E-01  -8.126E-02  -1.225E-04
    0.031   1.436E+00   9.592E-01  -6.276E-02  -2.714E+00   1.400E-01  -3.728E+00   2.343E-01  -5.430E-01  -6.448E-02  -3.230E-05
    0.025   1.522E+00   9.597E-01  -6.351E-02  -2.813E+00   1.458E-01  -3.654E+00   2.362E-01  -6.544E-01  -5.500E-02  -4.848E-05
    pga     9.069E-01   9.830E-01  -6.595E-02  -2.698E+00   1.594E-01  -2.795E+00   2.120E-01  -3.011E-01  -6.532E-02  -4.484E-04
    pgv    -1.442E+00   9.909E-01  -5.848E-02  -2.701E+00   2.155E-01  -2.436E+00   2.659E-01   8.479E-02  -6.927E-02  -3.734E-04
    """)

    #: Coefficients for NEHRP BC boundary (Vs30 = 760 m/s), table 9, pag 2202
    #: coefficient values taken from Fortran implementation of Dave Boore
    #: (higher precision than in the paper)
    COEFFS_BC = CoeffsTable(sa_damping=5, table="""\
    IMT     c1          c2          c3          c4          c5          c6          c7          c8          c9         c10
    5.000  -4.852E+00   1.580E+00  -8.066E-02  -2.530E+00   2.216E-01  -1.426E+00   1.361E-01   6.340E-01  -1.413E-01  -1.608E-04
    4.000  -5.256E+00   1.787E+00  -9.785E-02  -2.435E+00   2.068E-01  -1.307E+00   1.210E-01   7.340E-01  -1.560E-01  -1.959E-04
    3.125  -5.590E+00   1.972E+00  -1.136E-01  -2.331E+00   1.908E-01  -1.204E+00   1.099E-01   8.449E-01  -1.723E-01  -2.452E-04
    2.500  -5.800E+00   2.126E+00  -1.278E-01  -2.257E+00   1.790E-01  -1.123E+00   9.539E-02   8.911E-01  -1.797E-01  -2.601E-04
    2.000  -5.853E+00   2.233E+00  -1.385E-01  -2.195E+00   1.688E-01  -1.037E+00   8.002E-02   8.666E-01  -1.790E-01  -2.860E-04
    1.587  -5.754E+00   2.287E+00  -1.450E-01  -2.131E+00   1.582E-01  -9.568E-01   6.762E-02   8.670E-01  -1.789E-01  -3.429E-04
    1.250  -5.489E+00   2.289E+00  -1.476E-01  -2.081E+00   1.501E-01  -9.000E-01   5.794E-02   8.208E-01  -1.719E-01  -4.070E-04
    1.000  -5.058E+00   2.233E+00  -1.454E-01  -2.030E+00   1.408E-01  -8.744E-01   5.412E-02   7.922E-01  -1.697E-01  -4.886E-04
    0.794  -4.446E+00   2.119E+00  -1.387E-01  -2.009E+00   1.356E-01  -8.576E-01   4.976E-02   7.084E-01  -1.589E-01  -5.751E-04
    0.629  -3.748E+00   1.973E+00  -1.294E-01  -1.997E+00   1.313E-01  -8.417E-01   4.820E-02   6.772E-01  -1.557E-01  -6.763E-04
    0.500  -3.007E+00   1.803E+00  -1.178E-01  -1.982E+00   1.274E-01  -8.466E-01   4.698E-02   6.670E-01  -1.546E-01  -7.676E-04
    0.397  -2.281E+00   1.629E+00  -1.054E-01  -1.967E+00   1.227E-01  -8.880E-01   5.033E-02   6.839E-01  -1.582E-01  -8.587E-04
    0.315  -1.560E+00   1.455E+00  -9.312E-02  -1.977E+00   1.209E-01  -9.466E-01   5.576E-02   6.499E-01  -1.558E-01  -9.552E-04
    0.251  -8.756E-01   1.293E+00  -8.193E-02  -2.014E+00   1.226E-01  -1.027E+00   6.341E-02   5.808E-01  -1.491E-01  -1.053E-03
    0.199  -3.056E-01   1.156E+00  -7.211E-02  -2.038E+00   1.220E-01  -1.147E+00   7.375E-02   5.082E-01  -1.430E-01  -1.140E-03
    0.158   1.194E-01   1.057E+00  -6.473E-02  -2.054E+00   1.190E-01  -1.355E+00   9.160E-02   5.164E-01  -1.503E-01  -1.178E-03
    0.125   5.356E-01   9.647E-01  -5.835E-02  -2.110E+00   1.205E-01  -1.672E+00   1.156E-01   3.433E-01  -1.322E-01  -1.130E-03
    0.100   7.818E-01   9.235E-01  -5.555E-02  -2.165E+00   1.191E-01  -2.097E+00   1.483E-01   2.847E-01  -1.319E-01  -9.897E-04
    0.079   9.667E-01   9.033E-01  -5.476E-02  -2.249E+00   1.215E-01  -2.530E+00   1.775E-01   1.001E-01  -1.147E-01  -7.724E-04
    0.063   1.109E+00   8.875E-01  -5.386E-02  -2.334E+00   1.229E-01  -2.881E+00   2.007E-01  -3.189E-02  -1.069E-01  -5.483E-04
    0.050   1.209E+00   8.830E-01  -5.441E-02  -2.440E+00   1.295E-01  -3.035E+00   2.133E-01  -2.098E-01  -8.997E-02  -4.145E-04
    0.040   1.261E+00   8.789E-01  -5.515E-02  -2.536E+00   1.388E-01  -2.994E+00   2.158E-01  -3.908E-01  -6.746E-02  -3.881E-04
    0.031   1.191E+00   8.884E-01  -5.642E-02  -2.577E+00   1.451E-01  -2.840E+00   2.121E-01  -4.370E-01  -5.866E-02  -4.329E-04
    0.025   1.052E+00   9.030E-01  -5.768E-02  -2.571E+00   1.483E-01  -2.652E+00   2.065E-01  -4.084E-01  -5.769E-02  -5.122E-04
    pga     5.233E-01   9.686E-01  -6.196E-02  -2.439E+00   1.465E-01  -2.335E+00   1.912E-01  -8.695E-02  -8.285E-02  -6.304E-04
    pgv    -1.662E+00   1.050E+00  -6.035E-02  -2.496E+00   1.840E-01  -2.301E+00   2.500E-01   1.268E-01  -8.704E-02  -4.266E-04
    """)

    #: IMT-independent coefficients. std_total is the total standard deviation,
    #: see Table 6, pag 2192 and Table 9, pag 2202. R0, R1, R2 are coefficients
    #: required for mean calculation - see equation (5) pag 2191. v1, v2, Vref
    #: are coefficients required for soil response calculation, see table 8,
    #: p. 2201
    COEFFS_IMT_INDEPENDENT = {
        # the std is converted from base 10 to base e
        'std_total': np.log(10 ** 0.30),
        'R0': 10.0,
        'R1': 70.0,
        'R2': 140.0,
        'v1': 180.0,
        'v2': 300.0,
        'Vref': 760.0
    }

    STRESS_COEFFS = CoeffsTable(sa_damping=5, table="""\
    IMT    delta  M1    Mh
    pga    0.15   0.50  5.50
    0.025  0.15   0.00  5.00
    0.031  0.15   0.00  5.00
    0.04   0.15   0.00  5.00
    0.05   0.15   0.00  5.00
    0.063  0.15   0.17  5.17
    0.079  0.15   0.34  5.34
    0.1    0.15   0.50  5.50
    0.126  0.15   1.15  5.67
    0.158  0.15   1.85  5.84
    0.199  0.15   2.50  6.00
    0.251  0.15   2.90  6.12
    0.315  0.15   3.30  6.25
    0.397  0.15   3.65  6.37
    0.5    0.15   4.00  6.50
    0.629  0.15   4.17  6.70
    0.794  0.15   4.34  6.95
    1.00   0.15   4.50  7.20
    1.25   0.15   4.67  7.45
    1.587  0.15   4.84  7.70
    2.0    0.15   5.00  8.00
    2.5    0.15   5.25  8.12
    3.125  0.15   5.50  8.25
    4.0    0.15   5.75  8.37
    5.0    0.15   6.00  8.50
    pgv    0.11   2.00  5.50
    """)

class AtkinsonBoore2006MblgAB1987bar140NSHMP2008(AtkinsonBoore2006):
    """
    Implements GMPE developed by Gail M. Atkinson and David M. Boore and
    published as "Earthquake Ground-Motion Prediction Equations for Eastern
    North America" (2006, Bulletin of the Seismological Society of America,
    Volume 96, No. 6, pages 2181-2205) as utilized by the National Seismic
    Hazard Mapping Project (NSHMP) for the 2008 central and eastern US model.

    The class replicates the algorithm as coded in ``subroutine getAB06``
    in ``hazgridXnga2.f`` Fortran code available at:
    http://earthquake.usgs.gov/hazards/products/conterminous/2008/software/

    The class implement the equation for static stress drop equal to 140 bar.

    The class assumes rupture magnitude to be in Mblg scale (given that
    MFDs for central and eastern US are given in this scale). Therefore Mblg
    is converted to Mw by using Atkinson and Boore 1987 conversion equation.

    Mean value is clipped at 1.5 g for PGA and 3.0 g for SA with periods in
    range (0.02, 0.55) s.
    """
    def get_mean_and_stddevs(self, sites, rup, dists, imt, stddev_types):
        """
        See :meth:`superclass method
        <.base.GroundShakingIntensityModel.get_mean_and_stddevs>`
        for spec of input and result values.
        """
        mag = self._convert_magnitude(rup.mag)

        mean = self._get_mean(sites.vs30, mag, dists.rrup, imt, scale_fac=0)
        stddevs = self._get_stddevs(stddev_types, num_sites=sites.vs30.size)

        mean = clip_mean(imt, mean)

        return mean, stddevs

    def _convert_magnitude(self, mag):
        """
        Convert magnitude from Mblg to Mw using Atkinson and Boore 1987
        equation
        """
        return mblg_to_mw_atkinson_boore_87(mag)

class AtkinsonBoore2006MblgJ1996bar140NSHMP2008(
        AtkinsonBoore2006MblgAB1987bar140NSHMP2008):
    """
    Extend :class:`AtkinsonBoore2006MblgAB1987bar140NSHMP2008` but uses
    Johnston 1996 equation to convert from Mblg to Mw
    """
    def _convert_magnitude(self, mag):
        """
        Convert magnitude from Mblg to Mw using Johnston 1996 equation
        """
        return mblg_to_mw_johnston_96(mag)

class AtkinsonBoore2006Mwbar140NSHMP2008(
        AtkinsonBoore2006MblgAB1987bar140NSHMP2008):
    """
    Extend :class:`AtkinsonBoore2006MblgAB1987bar140NSHMP2008` but assumes
    magnitude to be in Mw scale and thefore no conversion is applied
    """
    def _convert_magnitude(self, mag):
        """
        Return magnitude value unchanged
        """
        return mag

class AtkinsonBoore2006MblgAB1987bar200NSHMP2008(AtkinsonBoore2006):
    """
    Same as :class:`AtkinsonBoore2006MblgAB1987bar140NSHMP2008` but with
    adjustment for 200 bar stress drop
    """
    def get_mean_and_stddevs(self, sites, rup, dists, imt, stddev_types):
        """
        See :meth:`superclass method
        <.base.GroundShakingIntensityModel.get_mean_and_stddevs>`
        for spec of input and result values.
        """
        mag = self._convert_magnitude(rup.mag)

        # stress drop scaling factor defined in subroutine getAB06
        mean = self._get_mean(
            sites.vs30, mag, dists.rrup, imt, scale_fac=0.5146
        )
        stddevs = self._get_stddevs(stddev_types, num_sites=sites.vs30.size)

        mean = clip_mean(imt, mean)

        return mean, stddevs

    def _convert_magnitude(self, mag):
        """
        Convert magnitude from Mblg to Mw using Atkinson and Boore 1987
        equation
        """
        return mblg_to_mw_atkinson_boore_87(mag)

class AtkinsonBoore2006MblgJ1996bar200NSHMP2008(
        AtkinsonBoore2006MblgAB1987bar200NSHMP2008):
    """
    Extend :class:`AtkinsonBoore2006MblgAB1987bar200NSHMP2008` but uses
    Johnston 1996 equation to convert from Mblg to Mw
    """
    def _convert_magnitude(self, mag):
        """
        Convert magnitude from Mblg to Mw using Johnston 1996 equation
        """
        return mblg_to_mw_johnston_96(mag)

class AtkinsonBoore2006Mwbar200NSHMP2008(
        AtkinsonBoore2006MblgAB1987bar200NSHMP2008):
    """
    Extend :class:`AtkinsonBoore2006MblgAB1987bar200NSHMP2008` but assumes
    magnitude to be in Mw scale therefore no conversion is applied
    """
    def _convert_magnitude(self, mag):
        """
        Return magnitude value unchanged
        """
        return mag

class AtkinsonBoore2006Modified2011(AtkinsonBoore2006):
    """
    This GMPE modifies the original implementation of :class:
    `AtkinsonBoore2006` with the magnitude dependent stress-drop scaling
    factor proposed in Atkinson & Boore (2011)
    Atkinson, G. A. and Boore D. M. (2011) Modifications to Existing
    Ground-Motion Prediciton Equations in Light of New Data. Bulletin of the
    Seismological Society of America, 101(3), 1121 - 1135
    """
    def get_mean_and_stddevs(self, sites, rup, dists, imt, stddev_types):
        """
        See :meth:`superclass method
        <.base.GroundShakingIntensityModel.get_mean_and_stddevs>`
        for spec of input and result values.
        """
        # Stress drop scaling factor is now a property of magnitude
        scale_fac = self._get_stress_drop_scaling_factor(rup.mag)
        mean = self._get_mean(sites.vs30, rup.mag, dists.rrup, imt, scale_fac)
        stddevs = self._get_stddevs(stddev_types, num_sites=sites.vs30.size)

        return mean, stddevs

    def _get_stress_drop_scaling_factor(self, magnitude):
        """
        Returns the magnitude dependent stress drop scaling factor defined in
        equation 6 (page 1128) of Atkinson & Boore (2011)
        """
        stress_drop = 10.0 ** (3.45 - 0.2 * magnitude)
        cap = 10.0 ** (3.45 - 0.2 * 5.0)
        if stress_drop > cap:
            stress_drop = cap
        return log10(stress_drop / 140.0) / log10(2.0)