import pytools.timetools as tt
import pytools.nctools as nct
import numpy as np
import netCDF4 as nc
import os


def interpolate(varName, des, srcs, maxDeltaHour=12, forced=False):
    print(f'from={[
        tt.float2format(src.getValidTime(), "%y/%m/%d_%H")
        for src in srcs
    ]}, to={tt.float2format(des.getValidTime(), "%y/%m/%d_%H")}', end=' ')

    desPath = des.getPath('patched')
    if os.path.exists(des.getPath('patched')) and not forced:
        print("des file already exists.", end=' ')
        print(des.getPath('patched'))
        return

    #
    # ---- check if the src paths exists
    for src in srcs:
        if not os.path.exists(src.getPath('source')):
            print(f'src path not found: {src.getPath("source")}')
            return

    #
    # ---- check delta hours to interpolate
    deltas = np.array([
        24 * np.abs(des.getValidTime() - src.getValidTime())
        for src in srcs 
    ])

    if any([delta > maxDeltaHour for delta in deltas]):
        print(f'deltas too big ({deltas} > max={maxDeltaHour})')
        return

    print(f'hours={deltas}')

    #
    # ---- read data and interpolate 
    datas = np.array([src.read() for src in srcs])
    times = np.array([src.getValidTime() for src in srcs])
    newTime = des.getValidTime()
    newData = datas[0, :] \
        + (datas[1, :] - datas[0, :]) \
        / (times[1] - times[0]) \
        * (newTime - times[0]) 

    #
    # ---- save data
    desDir = os.path.dirname(desPath)
    if not os.path.exists(desDir):
        os.system(f'mkdir -p {desDir}')

    if not os.path.exists(desPath):
        os.system(f'cp {srcs[0].getPath("source")} {desPath}')

    with nc.Dataset(desPath, 'a') as h:
        h[getNcVarName(varName)][:] = newData
    

class Forecast:
    def __init__(self, varName, initTime, leadHour):
        self.initTime = initTime
        self.leadHour = leadHour
        self.varName = varName

    def getValidTime(self):
        return self.initTime + self.leadHour/24
    
    def getPath(self, mode='source'):
        if mode not in ['patched', 'source',]:
            raise ValueError(f'unrecognized {mode=}')
        return getPath(self.varName, self.initTime, self.leadHour, mode) 

    def read(self, mode='source'):
        return nct.read(self.getPath(mode), getNcVarName(self.varName))
       
 
def getPath(varName, initTime, leadHour, mode):
    return tt.float2format(
        initTime,
        f'../data/{mode}/%Y/{varName}/%m/hycom_{varName}_%Y%m%d%H_t{leadHour:03d}.nc'
    )


def validTime2InitLead(validTime):
    year, month, day = tt.float2ymd(validTime)
    hour = tt.hour(validTime)
    if hour >= 12:
        initTime, lead = tt.ymd2float(year, month, day, 12), hour - 12
    else:
        initTime, lead = tt.ymd2float(year, month, day, 12) - 1, hour + 12
    return initTime, lead


def getNcVarName(varName):
    if varName == 'ssh':
        return 'surf_el'
    else:
        raise NotImplementedError(f'unrecognized {varName=}')