
# distutils: libraries = STD_LIBS
"""
Simple readers for fortran unformatted data, specifically for the Tiger code.



"""


import numpy as np

cimport cython
cimport numpy as np
from libc.stdio cimport FILE, fclose, fopen

#cdef inline int imax(int i0, int i1):
    #if i0 > i1: return i0
    #return i1

cdef extern from "endian_swap.h":
    void FIX_SHORT( unsigned short )
    void FIX_LONG( unsigned )
    void FIX_FLOAT( float )

cdef extern from "platform_dep.h":
    void *alloca(size_t)

cdef extern from "stdio.h":
    cdef int SEEK_SET
    cdef int SEEK_CUR
    cdef int SEEK_END
    int fseek(FILE *stream, long offset, int whence)
    size_t fread(void *ptr, size_t size, size_t nmemb, FILE *stream)
    long ftell(FILE *stream)
    char *fgets(char *s, int size, FILE *stream)

@cython.boundscheck(False)
@cython.wraparound(False)
def read_and_seek(char *filename, np.int64_t offset1,
                  np.int64_t offset2, np.ndarray buffer, int bytes):
    cdef FILE *f = fopen(filename, "rb")
    cdef void *buf = <void *> buffer.data
    cdef char line[1024]
    cdef size_t n = 1023
    fseek(f, offset1, SEEK_SET)
    fgets(line, n, f)
    fseek(f, offset2, SEEK_CUR)
    fread(buf, 1, bytes, f)
    fclose(f)

def count_art_octs(char *fn, long offset,
                   int min_level, int max_level,
                   int nhydro_vars,
                   level_info):
    cdef int nchild = 8
    cdef int next_record = -1, nLevel = -1
    cdef int dummy_records[9]
    cdef int readin = -1
    cdef FILE *f = fopen(fn, "rb")
    fseek(f, offset, SEEK_SET)
    for _ in range(min_level + 1, max_level + 1):
        fread(dummy_records, sizeof(int), 2, f)
        fread(&nLevel, sizeof(int), 1, f); FIX_LONG(nLevel)
        print(level_info)
        level_info.append(nLevel)
        fread(dummy_records, sizeof(int), 2, f)
        fread(&next_record, sizeof(int), 1, f); FIX_LONG(next_record)
        print("Record size is:", next_record)
        # Offset for one record header we just read
        next_record = (nLevel * (next_record + 2*sizeof(int))) - sizeof(int)
        fseek(f, next_record, SEEK_CUR)
        # Now we skip the second section
        fread(&readin, sizeof(int), 1, f); FIX_LONG(readin)
        nhydro_vars = next_record//4-2-3 #nhvar in daniel's code
        #record length is normally 2 pad bytes, 8 + 2 hvars (the 2 is nchem)
        # and then 3 vars, but we can find nhvars only here and not in other
        # file headers
        next_record = (2*sizeof(int) + readin) * (nLevel * nchild)
        next_record -= sizeof(int)
        fseek(f, next_record, SEEK_CUR)
    print("nhvars",nhydro_vars)
    fclose(f)

def read_art_tree(char *fn, long offset,
                  int min_level, int max_level,
                  np.ndarray[np.int64_t, ndim=2] oct_indices,
                  np.ndarray[np.int64_t, ndim=1] oct_levels,
                  np.ndarray[np.int64_t, ndim=2] oct_info):
    #             np.ndarray[np.int64_t, ndim=1] oct_mask,
    #             np.ndarray[np.int64_t, ndim=1] oct_parents,

    # This accepts the filename of the ART header and an integer offset that
    # points to the start of the record *following* the reading of iOctFree and
    # nOct.  For those following along at home, we only need to read:
    #   iOctPr, iOctLv
    cdef int nchild = 8
    cdef int iOct, nLevel, ic1
    cdef np.int64_t next_record = -1
    cdef long long child_record
    cdef int iOctPs[3]
    cdef np.int64_t dummy_records[9]
    cdef int readin = -1
    cdef FILE *f = fopen(fn, "rb")
    fseek(f, offset, SEEK_SET)
    cdef int Level = -1
    cdef int * iNOLL = <int *> alloca(sizeof(int)*(max_level-min_level+1))
    cdef int * iHOLL = <int *> alloca(sizeof(int)*(max_level-min_level+1))
    cdef int iOctMax = 0
    level_offsets = [0]
    for _ in range(min_level + 1, max_level + 1):
        fread(&readin, sizeof(int), 1, f); FIX_LONG(readin)
        fread(&Level, sizeof(int), 1, f); FIX_LONG(Level)
        fread(&iNOLL[Level], sizeof(int), 1, f); FIX_LONG(iNOLL[Level])
        fread(&iHOLL[Level], sizeof(int), 1, f); FIX_LONG(iHOLL[Level])
        fread(&readin, sizeof(int), 1, f); FIX_LONG(readin)
        iOct = iHOLL[Level] - 1
        nLevel = iNOLL[Level]
        #print("Reading Hierarchy for Level", Lev, Level, nLevel, iOct)
        #print(ftell(f))
        for ic1 in range(nLevel):
            iOctMax = max(iOctMax, iOct)
            #print(readin, iOct, nLevel, sizeof(int))
            next_record = ftell(f)
            fread(&readin, sizeof(int), 1, f); FIX_LONG(readin)
            assert readin==52
            next_record += readin + sizeof(int)
            fread(iOctPs, sizeof(int), 3, f)
            FIX_LONG(iOctPs[0]); FIX_LONG(iOctPs[1]); FIX_LONG(iOctPs[2])
            oct_indices[iOct, 0] = iOctPs[0]
            oct_indices[iOct, 1] = iOctPs[1]
            oct_indices[iOct, 2] = iOctPs[2]
            oct_info[iOct, 1] = ic1
            #grid_info[iOct, 2] = iOctPr # we don't seem to need this
            fread(dummy_records, sizeof(int), 6, f) # skip Nb
            fread(&readin, sizeof(int), 1, f); FIX_LONG(readin)
            #oct_parents[iOct] = readin - 1
            fread(&readin, sizeof(int), 1, f); FIX_LONG(readin)
            oct_levels[iOct] = readin
            fread(&iOct, sizeof(int), 1, f); FIX_LONG(iOct)
            iOct -= 1
            assert next_record > 0
            fseek(f, next_record, SEEK_SET)
            fread(&readin, sizeof(int), 1, f); FIX_LONG(readin)
            assert readin==52

        level_offsets.append(ftell(f))

        #skip over the hydro variables
        #find the length of one child section
        #print('measuring child record ',)
        fread(&next_record, sizeof(int), 1, f)
        #print(next_record,)
        FIX_LONG(next_record)
        #print(next_record)
        fseek(f,ftell(f)-sizeof(int),SEEK_SET) #rewind
        #This is a sloppy fix; next_record is 64bit
        #and I don't think FIX_LONG(next_record) is working
        #correctly for 64bits
        if next_record > 4294967296L:
            next_record -= 4294967296L
        assert next_record == 56

        #find the length of all of the children section
        child_record = ftell(f) +  (next_record+2*sizeof(int))*nLevel*nchild
        #print('Skipping over hydro vars', ftell(f), child_record)
        fseek(f, child_record, SEEK_SET)

        # for ic1 in range(nLevel * nchild):
        #     fread(&next_record, sizeof(int), 1, f); FIX_LONG(next_record)
        #     fread(&idc, sizeof(int), 1, f); FIX_LONG(idc); idc -= 1 + (128**3)
        #     fread(&cm, sizeof(int), 1, f); FIX_LONG(cm)
        #     #if cm == 0: oct_mask[idc] = 1
        #     #else: total_masked += 1
        #     assert next_record > 0
        #     fseek(f, next_record - sizeof(int), SEEK_CUR)
    fclose(f)
    return level_offsets

def read_art_root_vars(char *fn, long root_grid_offset,
                    int nhydro_vars, int nx, int ny, int nz,
                    int ix, int iy, int iz, fields, var):

    cdef FILE *f = fopen(fn, "rb")
    cdef int j,l, cell_record_size = nhydro_vars * sizeof(float)
    cdef float temp = -1
    l=0
    fseek(f, root_grid_offset, SEEK_SET)
    # Now we seet out the cell we want
    cdef int my_offset = (((iz * ny) + iy) * nx + ix)
    #print(cell_record_size, my_offset, ftell(f))
    fseek(f, cell_record_size * my_offset, SEEK_CUR)
    #(((C)*GridDimension[1]+(B))*GridDimension[0]+A)
    for j in range(nhydro_vars):
        fread(&temp, sizeof(float), 1, f)
        if j in fields:
            FIX_FLOAT(temp)
            var[l]=temp
            l+=1
    fclose(f)

cdef void read_art_vars(FILE *f,
                    int min_level, int max_level, int nhydro_vars,
                    int grid_level,long grid_id,long child_offset,
                    fields,
                    np.ndarray[np.int64_t, ndim=1] level_offsets,
                    var):
    # nhydro_vars is the number of columns- 3 (adjusting for vars)
    # this is normally 10=(8+2chem species)
    cdef int record_size = 2+1+1+nhydro_vars+2
    cdef float temp = -1.0
    cdef float varpad[2]
    cdef int new_padding = -1
    cdef int padding[3]
    cdef long offset = 8*grid_id*record_size*sizeof(float)
    fseek(f, level_offsets[grid_level] + offset, SEEK_SET)
    for j in range(8): #iterate over the children
        l = 0
        fread(padding, sizeof(int), 3, f); FIX_LONG(padding[0])
        #print("Record Size", padding[0])
        # This should be replaced by an fread of nhydro_vars length
        for k in range(nhydro_vars): #iterate over the record
            fread(&temp, sizeof(float), 1, f); FIX_FLOAT(temp)
            #print(k, temp)
            if k in fields:
                var[j,l] = temp
                l += 1
        fread(varpad, sizeof(float), 2, f)
        fread(&new_padding, sizeof(int), 1, f); FIX_LONG(new_padding)
        assert(padding[0] == new_padding)

@cython.cdivision(True)
@cython.boundscheck(False)
@cython.wraparound(False)
def read_art_grid(int varindex,
              np.ndarray[np.int64_t, ndim=1] start_index,
              np.ndarray[np.int32_t, ndim=1] grid_dims,
              np.ndarray[np.float32_t, ndim=3] data,
              np.ndarray[np.uint8_t, ndim=3] filled,
              np.ndarray[np.float32_t, ndim=2] level_data,
              int level, int ref_factor,
              component_grid_info):
    cdef int gi, i, j, k, grid_id
    cdef int ir, jr, kr
    cdef int offi, offj, offk, odind
    cdef np.int64_t di, dj, dk
    cdef np.ndarray[np.int64_t, ndim=1] ogrid_info
    cdef np.ndarray[np.int64_t, ndim=1] og_start_index
    cdef np.float64_t temp_data
    cdef np.int64_t end_index[3]
    cdef int kr_offset, jr_offset, ir_offset
    cdef int to_fill = 0
    # Note that indexing into a cell is:
    #   (k*2 + j)*2 + i
    for i in range(3):
        end_index[i] = start_index[i] + grid_dims[i]
    for gi in range(len(component_grid_info)):
        ogrid_info = component_grid_info[gi]
        grid_id = ogrid_info[1]
        og_start_index = ogrid_info[3:6] #the oct left edge
        for i in range(2*ref_factor):
            di = i + og_start_index[0] * ref_factor
            if di < start_index[0] or di >= end_index[0]: continue
            ir = <int> (i / ref_factor)
            for j in range(2 * ref_factor):
                dj = j + og_start_index[1] * ref_factor
                if dj < start_index[1] or dj >= end_index[1]: continue
                jr = <int> (j / ref_factor)
                for k in range(2 * ref_factor):
                    dk = k + og_start_index[2] * ref_factor
                    if dk < start_index[2] or dk >= end_index[2]: continue
                    kr = <int> (k / ref_factor)
                    offi = di - start_index[0]
                    offj = dj - start_index[1]
                    offk = dk - start_index[2]
                    #print(offi, filled.shape[0],)
                    #print(offj, filled.shape[1],)
                    #print(offk, filled.shape[2])
                    if filled[offi, offj, offk] == 1: continue
                    if level > 0:
                        odind = (kr*2 + jr)*2 + ir
                        # Replace with an ART-specific reader
                        #temp_data = local_hydro_data.m_var_array[
                        #        level][8*offset + odind]
                        temp_data = level_data[varindex, 8*grid_id + odind]
                    else:
                        kr_offset = kr + <int> (start_index[0] / ref_factor)
                        jr_offset = jr + <int> (start_index[1] / ref_factor)
                        ir_offset = ir + <int> (start_index[2] / ref_factor)
                        odind = (kr_offset * grid_dims[0] + jr_offset)*grid_dims[1] + ir_offset
                        temp_data = level_data[varindex, odind]
                    data[offi, offj, offk] = temp_data
                    filled[offi, offj, offk] = 1
                    to_fill += 1
    return to_fill

@cython.cdivision(True)
@cython.boundscheck(True)
@cython.wraparound(False)
def fill_child_mask(np.ndarray[np.int64_t, ndim=2] file_locations,
                    np.ndarray[np.int64_t, ndim=1] grid_le,
                    np.ndarray[np.uint8_t, ndim=4] art_child_masks,
                    np.ndarray[np.uint8_t, ndim=3] child_mask):

    #loop over file_locations, for each row extracting the index & LE
    #of the oct we will pull pull from art_child_masks
    #then use the art_child_masks info to fill in child_mask
    cdef int i,ioct,x,y,z
    cdef int nocts = file_locations.shape[0]
    cdef int lex,ley,lez
    for i in range(nocts):
        ioct = file_locations[i,1] #from fortran to python indexing?
        lex = file_locations[i,3] - grid_le[0] #the oct left edge x
        ley = file_locations[i,4] - grid_le[1]
        lez = file_locations[i,5] - grid_le[2]
        for x in range(2):
            for y in range(2):
                for z in range(2):
                    child_mask[lex+x,ley+y,lez+z] = art_child_masks[ioct,x,y,z]
