import csv
import sys

import numpy as np
import scipy.interpolate as interpolate
import scipy.stats as stats

def find_calibrator_px(shard_dict, config):

    """
    Takes the calibrator wavelengths and finds their corresponding px.

    Takes the wavelengths of the water calibration lines from the config file
    and returns the pixel closest to each wavelength and its shard. If a
    wavelength does not fall within any shard's wavelength range, raises a
    warning and excludes that wavelength. If no wavelength has a pixel closest
    to it, throws an error.
    """

    # 2) Find calibrator pixel for each calibration wavelength
    calibrator_px = []
    for wv in config["calibrators"]:

        found_match = False

        for shardloc, shard in shard_dict.items():
            lin_x = next(iter(shard.spectra.values())).lin_x
            if lin_x[0] < wv and wv < lin_x[-1]:
                px = (np.abs(lin_x - wv)).argmin()
                calibrator_px.append((shardloc, px))
                found_match = True
                break

        if not found_match:
            sys.stderr.write("No shard contained calibrator: {}A\n".format(wv))

    # 3) Make sure that there is at least one calibrator pixel
    if len(calibrator_px) == 0:
        raise Exception("No calibrators were contained by a shard.")

    return calibrator_px
            

def generate_calibrators(shard_dict, calibrator_px):
    """
    Generates the water and airmass calibrators.

    Lists i) the average height of the water calibration pixels in each spectra
    and ii) their airmasses, and records the order they are listed in.
    """

    # 1) Generate list of spectrum airmasses, and records the spectra's 
    # ordering.
    z_calibrator = []
    f_order = []
    # NOTE: We need to pick an order which contains every spectra we use:
    # not an order which has some spectra filtered out of it. I'd like to
    # rewrite this code properly, but the project specifications keep 
    # changing and I have no time. Order 72 is an order which shouldn't
    # have an spectra in it. I apologize to any one who happens to want to 
    # maintain this code.
    good_order = 72
    for filename, spectrum in shard_dict[good_order].spectra.items():
        z_calibrator.append(spectrum.z)
        f_order.append(filename)

    # 2) Generate list of the average height of the water calibration pixels
    # in each spectrum in the recorded order.
    print("calibrator_px:{}".format(calibrator_px))
    w_calibrator = None
    for order, px in calibrator_px:
        shard = shard_dict[order]
        px_heights = []
        for filename in f_order:
            spectrum = shard.spectra[filename]
            px_heights.append(spectrum.log_y[px])
        if w_calibrator is None:
            w_calibrator = np.array(px_heights)
        else:
            w_calibrator += px_heights

    w_calibrator /= len(calibrator_px)

    return (w_calibrator, z_calibrator, f_order)


def compute_PCC_threshold(p_value, db_path):

    """
    Computes the threshold PCC from user submitted p value.
    
    Notes
    -----
    For small n (< 300), scipy's stats module is unreliable. Thus, a
    simulation was done across 100,000 trials to generate a culmulative
    histogram of PCCs generated by noise, and the xth percentile p value
    calculated by taking the xth percentile PCC in the simulation (see 
    Leet et al. Tracking Tellurics in High Resolution Spectra,
    I. A Linear Regression Model, ApJ, 2019.) This is a temporary fix and will
    be improved.
    """

    # 1) Read p_values and PCCs from the k threshold databse
    p_values, PCCs = [], []
    PVAL_I, PCC_I = 0, 1
    with open(db_path) as db:
        db_reader = csv.reader(db, delimiter=" ")
        are_headers_read = False
        for record in db_reader:
            if are_headers_read:
                p_values.append(float(record[PVAL_I]))
                PCCs.append(float(record[PCC_I]))
            else:
                are_headers_read = True

    # 2) Create mapping between p_values and PCCs by linear interpolation.
    p_value_2_PCC = interpolate.interp1d(p_values, PCCs)
    
    # 3) Use mapping to determine threshold p_value for as long as p_value
    # is between 0.1 and 0.00001
    if p_value > 0.1 or p_value < 0.00001:
        raise Exception("p_value in config outside of supported range 0.1-0.00001")
    k = p_value_2_PCC(p_value)
    return k
                

def flag_high_PCC_pixels(calibrators, k, shard_dict, config):

    """
    Flag pixels with high PCC with either the water calibrator or airmass.

    Computes the PCC of each pixel in each shard with the water calibrator
    and with airmass. Sets a pixel's airmass flag if its growth's PCC with
    airmass is above the p-value given in config, and a pixel's water 
    flag if its growth's PCC is above the p-value given in config.
    """

    for shard in shard_dict.values():
        flag_high_PCC_pixels_in_shard(calibrators, k, shard, config)

    
def flag_high_PCC_pixels_in_shard(calibrators, k, shard, config):

    """
    Flag pixels with high PCC with the water calibrator or airmass in shard.

    Worker function of flag_high_PCC_pixels (above)
    """

    w_calibrator, z_calibrator, f_order = calibrators
    
    # 0) Initialize px arrays
    shard.w_PCCs = np.zeros(shard.px_no)
    shard.w_tel = np.zeros(shard.px_no)
    shard.z_PCCs = np.zeros(shard.px_no)
    shard.z_tel = np.zeros(shard.px_no)
    shard.s_tel = np.zeros(shard.px_no)

    for i in range(shard.px_no):
        
        # 1) Order each pixel's spectrum intensities in the same order of
        # spectra as the calibrators.
        px_ws, px_zs, px_intensities = [], [], []
        for w, z, filename in zip(w_calibrator, z_calibrator, f_order):
            if filename in shard.spectra:
              spectrum = shard.spectra[filename]
              px_ws.append(w)
              px_zs.append(z)
              px_intensities.append(spectrum.log_y[i])

        # 2) Record the pixel intensity's PCC with water and flag pixels
        # with significant PCC at the config's p value.
        PCC = np.corrcoef(px_ws, px_intensities)[0][1]
        shard.w_PCCs[i] = PCC
        shard.w_tel[i] = PCC > k

        # 3) Record the pixel intensity's PCC with z and flag pixels
        # with significant PCC at the config's p value.
        PCC = np.corrcoef(px_zs, px_intensities)[0][1]
        shard.z_PCCs[i] = (-PCC)
        shard.z_tel[i] = (-PCC) > k

        # 4) If the average intensity is > saturation_threshold, flag
        # pixel as saturated
        shard.s_tel[i] = np.mean(px_intensities) < config["saturation_threshold"]
            
