import struct
import numpy as np
import pdb 
from scipy import interpolate
from pyproj import Transformer
import rioxarray  # This is required to add the .rio functionality
import xarray as xr
from pyproj import CRS

def slread( fname, Tstart, Tend, *args, **kwargs ):

    """
    Reads FDS slice file
    Based on slread.m by Simo Hostikka
    https://github.com/firemodels/fds/blob/master/Utilities/Matlab/scripts/slread.m
    (QQ,Time)=slread(fname,Tstart,Tend [,Nframes, gridskip, timeskip]);
      Tstart  is start time
      Tend    is end time
      Nframes is number of slice file frames (FDS default is 1000 between Tstart and Tend)
      gridskip is a skip rate for reading cells: =2 --> read every other cell, etc.
      timeskip is a skip rate for frames: =2 --> read every other frame, etc.
      QQ      contains the data
      Time    contains the time points
    """

    print(fname)
    print(Tstart)
    print(Tend)

    if len(args)==0:
        Nframes = 10000
    else:
        Nframes = args[0]
        
    gridskip = kwargs['gridskip'] if 'gridskip' in kwargs else 1
    timeskip = kwargs['timeskip'] if 'timeskip' in kwargs else 1
        
    f = open(fname,'rb')

    print('#######')
    f.read(4)
    Str1 = f.read(30)       # quantity
    print(Str1)
    f.read(8)
    Str2 = f.read(30)       # short name
    print(Str2)
    f.read(8)
    Str3 = f.read(30)       # units
    print(Str3)
    f.read(8)
    Indx = struct.unpack('6i',f.read(24))  # index bounds i,j,k
    print(Indx)
    f.read(4)

    # allocate arrays for QQ and Time

    Isize = Indx[1]-Indx[0]+1
    Jsize = Indx[3]-Indx[2]+1
    Ksize = Indx[5]-Indx[4]+1
    if Isize==1:
       M = Jsize
       N = Ksize
    elif Jsize==1:
       M = Isize
       N = Ksize
    else:
       M = Isize
       N = Jsize

    Nframes = max(1,Nframes)
    QQ = np.zeros((M,N))
    Time = np.zeros(Nframes+1)
    
    ii = np.arange(0,M,gridskip)
    jj = np.arange(0,N,gridskip)
    tt = np.arange(0,Nframes+1,timeskip)
    Qskip  = np.zeros((len(ii),len(jj),len(tt)))

    st = 0

    while Time[st] < Tstart:

        f.read(4)
        Time_list = struct.unpack('f',f.read(4))
        Time[st] = Time_list[0]
        f.read(8)
        for n in range(N):
            for m in range(M):
                QQ_list = struct.unpack('f',f.read(4))
                QQ[m,n] = QQ_list[0]
        f.read(4)

    while Time[st-1] < Tend:

        f.read(4)
        try: 
            Time_list = struct.unpack('f',f.read(4))
        except: 
            pdb.set_trace()
        Time[st] = Time_list[0]
        f.read(8)
        for n in range(N):
            for m in range(M):
                QQ_list = struct.unpack('f',f.read(4))
                QQ[m,n] = QQ_list[0]
        if st%timeskip==0:
            #print(st,int(st/timeskip))
            Qskip[:,:,int(st/timeskip)] = QQ[np.ix_(ii,jj)]
        f.read(4)
        st = st + 1
        if st>Nframes:
            break

    return(Qskip[:,:,:st],Time[:st])

def getTerrain(fname):
    f = open(fname,'rb')

    #zmin
    struct.unpack('i',f.read(4))
    zmin = struct.unpack('f',f.read(4))
    struct.unpack('i',f.read(4))

    #nx, ny
    struct.unpack('i',f.read(4))
    nx = struct.unpack('i',f.read(4))[0]
    ny = struct.unpack('i',f.read(4))[0]
    struct.unpack('i',f.read(4))

    #x
    lenx = struct.unpack('i',f.read(4))[0]
    x = np.array(struct.unpack('{:d}f'.format(lenx//4),f.read(lenx)))
    struct.unpack('i',f.read(4))

    #y
    leny=struct.unpack('i',f.read(4))[0]
    y=np.array(struct.unpack('{:d}f'.format(leny//4),f.read(leny)))
    struct.unpack('i',f.read(4))

    #z
    lenz = struct.unpack('i',f.read(4))[0]
    zz = np.array(struct.unpack('{:d}f'.format(lenz//4),f.read(lenz)))
    struct.unpack('i',f.read(4))

    f.close()
    

    zz = zz.reshape((nx,ny))

    return x,y,zz,zmin



nproc=40
inputDir = '/data/paugam/FDS/LS-elpontdeVilomara/practical4/3-fdsls-pdVHouse-canyon-LS1/'
simName = 'pdV'
crs_fds = CRS.from_string("WGS 84 / UTM zone 31N")
x0 = 408033.5
y0 = 4618433.3

xa = []
ya = []
zza = []
for imesh in range(nproc):
    fname = '{:s}/{:}_{:d}.ter'.format(inputDir,simName,imesh+1)
    x,y,zz,zmin = getTerrain(fname)

    xa.append(x)
    ya.append(y)
    zza.append(zz)

nx = [xx.shape[0]  for xx in xa]
ny = [yy.shape[0]  for yy in ya]

#Big domain # assume all mesh have same resolution
dx = np.diff(xa[0]).mean()
dy = np.diff(ya[0]).mean()

print('resolution = ',dx,dy)

xamin = min([xx.min() for xx in xa ]) 
xamax = max([xx.max() for xx in xa ]) 
yamin = min([yy.min() for yy in ya ]) 
yamax = max([yy.max() for yy in ya ]) 

xmin = np.round(xamin,3)
xmax = np.round(xamax,3)
ymin = np.round(yamin,3)
ymax = np.round(yamax,3)

xbig = np.arange(xmin,xmax,dx)
ybig = np.arange(ymin,ymax,dy)

grid_y, grid_x = np.meshgrid(ybig, xbig)

points = []
valuesTer = []
for imesh in range(nproc):
    for i in range(nx[imesh]):
        for j in range(ny[imesh]):
            points.append([xa[imesh][i],ya[imesh][j]])
            valuesTer.append(zza[imesh][i,j])
    

terrain = interpolate.griddata(points, valuesTer, (grid_x, grid_y), method='linear', fill_value=-999)
    
# Create the dataset in epsg:4326
transformer = Transformer.from_crs(crs_fds, "EPSG:4326", always_xy=True)

# Perform the transformation to UTM
lon, lat = transformer.transform(grid_x+x0, grid_y+y0)

lats = np.round(lat[0,:],10)
lons = np.round(lon[:,0],10)

ds = xr.Dataset({
      'terrainFDS': xr.DataArray(
          data = terrain.T,
          dims = ['latitude','longitude'],
          coords = {'latitude': lats, 'longitude': lons},
          attrs = {'long_name': 'random data', 'units': 'm', '_fillValue': -999}
          )

      },
      attrs = {'example_attr': 'this is a global attribute'}
      )
  
# save datasets
ds = ds.rio.write_crs("EPSG:4326", inplace=True)
ds.to_netcdf('{:}/{:}_terrain.nc'.format('./',simName),format='NETCDF4')