"""
Created on Feb. 2 2018

@author: kanazawa
"""
import numpy as np
import os
from scipy import interpolate
import matplotlib.pyplot as plt

import Common.FileReader as rd

def UpSamplingAndNormalizationForUnet(width,
                      input_size,
                      chromatoDict,            #[in]
                      chromatos,               #[in, out]
                      ):
    times = chromatoDict["time"]
    intensities = chromatoDict["intensity"] 
    assert times[-1] - times[0] <= width

    padTime = width - (times[-1] - times[0]) % width
    minRate = min(np.diff(times))
    padTimes = np.arange(times[-1] + minRate, times[-1] + padTime, minRate)
    times = np.append(times, padTimes)
    intensities = np.pad(intensities, [0, len(padTimes)], "edge")

    orgIntensities = intensities
    orgUpsamplingInts = []
    orgUpsamplingTimes = []

    startTime = times[0]
    endTime = times[-1]

    assert(min(intensities) >= 0)
    logIntensities = np.log(orgIntensities + 1.0) 
    diff1Intensities = np.diff(np.r_[orgIntensities, orgIntensities[-1]], n= 1) 
    diff2Intensities = np.diff(np.r_[orgIntensities[0], orgIntensities, orgIntensities[-1]], n= 2) 
    
    f1 = interpolate.interp1d(times, intensities, kind = "cubic")
    upsamplingTimes = np.linspace(startTime, endTime, input_size)
    upsamplingIntensities = f1(upsamplingTimes)
    assert not(np.any(np.isnan(upsamplingIntensities)))
    orgUpsamplingTimes.append(upsamplingTimes)
    orgUpsamplingInts.append(upsamplingIntensities)

    modifiedUpsamplingIntensities = upsamplingIntensities - min(upsamplingIntensities)
    if(0 != max(modifiedUpsamplingIntensities)):
        modifiedUpsamplingIntensities = modifiedUpsamplingIntensities / max(modifiedUpsamplingIntensities) * 2 - 1.0

    f1 = interpolate.interp1d(times, logIntensities, kind = "linear")
    upsamplingLogIntensities = f1(upsamplingTimes)
    assert not(np.any(np.isnan(upsamplingLogIntensities)))
    upsamplingLogIntensities -= min(upsamplingLogIntensities)
    if(0 != max(upsamplingLogIntensities)):
        upsamplingLogIntensities = upsamplingLogIntensities / max(upsamplingLogIntensities) * 2 - 1.0
    
    f1 = interpolate.interp1d(times, diff1Intensities, kind = "cubic")
    upsampling1DiffIntensities = f1(upsamplingTimes)
    assert not(np.any(np.isnan(upsampling1DiffIntensities)))
    upsampling1DiffIntensities -= min(upsampling1DiffIntensities)
    if(0 != max(upsampling1DiffIntensities)):
        upsampling1DiffIntensities = upsampling1DiffIntensities / max(upsampling1DiffIntensities) * 2 - 1.0

    f1 = interpolate.interp1d(times, diff2Intensities, kind = "cubic")
    upsampling2DiffIntensities = f1(upsamplingTimes)
    assert not(np.any(np.isnan(upsampling2DiffIntensities)))
    upsampling2DiffIntensities -= min(upsampling2DiffIntensities)
    if(0 != max(upsampling2DiffIntensities)):
        upsampling2DiffIntensities = upsampling2DiffIntensities / max(upsampling2DiffIntensities) * 2 - 1.0

    chromatos.append(np.asarray([modifiedUpsamplingIntensities,
                                    upsamplingLogIntensities,
                                    upsampling1DiffIntensities,
                                    upsampling2DiffIntensities]).reshape(4, input_size).astype(np.float32))    
        
    return orgUpsamplingTimes, orgUpsamplingInts
