from netCDF4 import Dataset
import numpy as np
import xarray as xr
import os
from wrf import getvar, interpline, CoordPair, xy_to_ll, ll_to_xy
import pandas as pd
from astropy.table import Table
from astropy.io import ascii

# List the months and years we want to plot in our time series
years_arr = np.full(11,-999,dtype=int)
years_arr = [1990,1991,1992,1993,1994,1995,1996,1997,1998,1999,2000]
numYEAR = len(years_arr)
years_str = ["1990","1991","1992","1993","1994","1995","1996","1997","1998","1999","2000"]

# Calculate how many points there will be in our time series (132 points)
days = 365
days_span = np.arange(0,days,1)

# Make array for ascii txt file
data = np.full([11,5,days],1e20,dtype=float)

# Here we set up our netcdf file that we want to output and create dimensions and variables
os.system("rm -rf t2_precip_rtty_1990.nc")
out = Dataset("t2_precip_rtty_1990.nc","w")
year = np.copy(years_arr)
out.createDimension('Year',None)
out.createDimension('Time',days)
YEARS = out.createVariable('YEARS',year.dtype,'Year')
out.variables['YEARS'][:]=year[:]
T2MAX = out.createVariable('T2MAX','f4',('Year','Time'))
PREC = out.createVariable('PREC','f4',('Year','Time'))

# Loop through the years so that we get each day of each year for our file.
i = 0
for years in years_str:
    print(years)

# Add the t2 maxmin daily file we just made
    a = Dataset("data/t2_maxmin_daily_"+years+"_rtty_1990.nc","r")
    c = Dataset("data/precip_daily_"+years+"_rtty_1990.nc","r")
    b = Dataset("data/met_em.d01.2004-12-16_12:00:00.nc","r")

# Here we run the function wrf_user_ll_to_ij.  We send in lat/lon and it returns the i,j
#   grid box that we care about.
    i_j = ll_to_xy(b,40.0274,-105.2705)  # Lat/lon for Boulder, CO
    locX = i_j[0]
    locY = i_j[1]

# Pull out the data for the grid box we care about.
    t2 = a.variables['T2MAX'][:,locY,locX]
    prec = c.variables['PREC_DAY'][:,locY,locX]

# Write out the data to the netcdf file
    out.variables['T2MAX'][i,:]=t2[:]
    out.variables['PREC'][i,:]=prec[:]

# Create the dates for each year and deal with leap years
    date = pd.date_range(start='1/1/'+years,end='12/31/'+years,freq='D')
    date = date[~((date.month == 2) & (date.day == 29))]

# Now put data into a table to be written out to a text file
    data[i,0,:] = date.year
    data[i,1,:] = date.month
    data[i,2,:] = date.day

    for day in days_span:
        data[i,3,day] = t2[day]
        data[i,4,day] = prec[day]
    data_tab = Table([data[i,0,:],data[i,1,:],data[i,2,:],data[i,3,:],data[i,4,:]],\
                  names=('Y','M','D','T2','PREC'))
    if i==0:
        ascii.write(data_tab,'t2_precip_rtty_1990.txt',format='no_header',\
                    formats={'Y': '%4d','M': '%2d','D': '%2d','T2': '%3.2f','PREC': '%3.2f'},\
                    overwrite=True)
    else:
        with open('t2_precip_rtty_1990.txt',mode='a') as f:
            data_tab.write(f,format='ascii.no_header',\
                        formats={'Y': '%4d','M': '%2d','D': '%2d','T2': '%3.2f','PREC': '%3.2f'})
    i = i+1
out.close()