import numpy as np
import os
import numpy as np
from .binSubfunctions import *

def character_indexes(string, match):
    indexes_of_match = []

    for index, character in enumerate(string):
        if character == match:
            indexes_of_match.append(index)
    return indexes_of_match

def GetVariableName(string):
    indexes_of_match = character_indexes(string, '"')
    return string[indexes_of_match[0] + 1 : indexes_of_match[1]]
class ProcData:
    def __init__(self, filename=None):
        self.filename = filename
        self.data = {}
        self.nse_read_bin_data = {}
        self.variable_names = ""

    def __get_variable_names_plt(self):
        with open(self.filename, "r") as file:
            file_data = file.read().split('\n')
        file.close()

        names = [GetVariableName(variable_name) for variable_name in file_data[1].split('=')[1].split(',')]
        self.variable_names = np.array(names, dtype=object)
    
    def __get_variable_names_bin(self):
        f = open(self.filename,'rb')
        _ = read_int(f)
        _ = read_int(f)
        ndims = read_int(f)
        data_type_size = read_int(f)

        if ndims == 1:
            self.variable_names = ['I']
        elif ndims == 2:
            self.variable_names = ['I', 'J']
        elif ndims == 3:
            self.variable_names = ['I', 'J', 'K']

        if data_type_size == 4:
            read_float = read_float32
            read_floats = read_floats32
            float_type = np.float32
        elif data_type_size == 8:
            read_float = read_float64
            read_floats = read_floats64
            float_type = np.float64
        else:
            raise Exception(f"Unknown data_type value: {data_type_size}. Must be either 4 or 8.")

        for _ in range(24):
            read_float(f)

        nx = read_int(f)
        ny = read_int(f)
        nz = read_int(f)

        # skip gap; reserved for future usage
        for _ in range(21):
            read_int(f)

        for _ in range(2):
            read_floats(f, num=nx)
            read_floats(f, num=ny)
            read_floats(f, num=nz)

        read_float(f)

        field_type = read_int(f)
        if field_type == 0:
            nvars = 1
        elif field_type == 1:
            nvars = ndims
        else:
            raise Exception(f"Unknown field_type value: {field_type}. Must be either 0 or 1")

        len_varnames = []
        for _ in range(nvars):
            len_varnames += [read_int(f)]

        varnames = []
        for nvar in range(nvars):
            varnames += [read_str(f, length=len_varnames[nvar])]
                    
        self.variable_names += varnames
        self.variable_names = np.array(self.variable_names, dtype=object)
    
    def __read_plt(self, ndim):
        with open(self.filename, "r") as file:
            file_data = file.read().split('\n')
        file.close()

        I = int(file_data[2].split(',')[0].split('=')[1])

        if ndim == 1:
            var_begin = 1
            J = 0
            K = 0
        elif ndim == 2:
            var_begin = 2
            J = int(file_data[2].split(',')[1].split('=')[1])
            K = 0
        elif ndim == 3:
            var_begin = 3
            J = int(file_data[2].split(',')[1].split('=')[1])
            K = int(file_data[2].split(',')[2].split('=')[1])

        file_data = np.loadtxt(self.filename, dtype = float, skiprows=3, unpack=False, ndmin=2, encoding='utf-8')
        file_data = file_data.T

        if self.var == None:
            column_names = [self.variable_names[i] for i in range(var_begin, len(self.variable_names))] 
        else:
            column_names = self.var
            
        if self.mval != None:
            for col, abs_val in zip(column_names, self.mval):
                check_variable_idx = np.where(self.variable_names == col)[0][0]
                idx = np.where(np.abs(file_data[check_variable_idx]) > abs_val)[0]
                file_data[check_variable_idx][idx] = np.nan

        self.data[self.variable_names[0]] = file_data[0][0:I]

        if ndim == 2:
            self.data[self.variable_names[1]] = file_data[1][: : I]
        elif ndim == 3:
            self.data[self.variable_names[1]] = file_data[1][: : I]
            self.data[self.variable_names[2]] = file_data[2][: : I * J]

        i = 0
        for name in self.variable_names[ndim : ]:
            if ndim == 1:
                self.data[name] = file_data[ndim + i]
            elif ndim == 2:
                self.data[name] = file_data[ndim + i].reshape( J, I )
            elif ndim == 3:
                self.data[name] = file_data[ndim + i].reshape( K, J, I )
            i += 1

        self.data['cx'] = I
        self.data['cy'] = J
        self.data['cz'] = K

    
    def __read_bin(self):
        f = open(self.filename,'rb')
        fid = read_int(f)
        grid_type = read_int(f)
        ndims = read_int(f)
        data_type_size = read_int(f)

        if data_type_size == 4:
            read_float = read_float32
            read_floats = read_floats32
            float_type = np.float32
        elif data_type_size == 8:
            read_float = read_float64
            read_floats = read_floats64
            float_type = np.float64
        else:
            raise Exception(f"Unknown data_type value: {data_type_size}. Must be either 4 or 8.")
        
        x = read_float(f)
        y = read_float(f)
        z = read_float(f)

        xdomainsize = read_float(f)
        ydomainsize = read_float(f)
        zdomainsize = read_float(f)

        # skip gap; reserved for future usage
        for _ in range(18):
            read_float(f)

        nx = read_int(f)
        ny = read_int(f)
        nz = read_int(f)

        if ndims < 3:
            nz = 0
        if ndims < 2:
            ny = 0

        dimsizes = ([nx, ny, nz])[0:ndims]

        gcx = read_int(f)
        gcy = read_int(f)
        gcz = read_int(f)

        # skip gap; reserved for future usage
        for _ in range(18):
            read_int(f)

        cx = read_floats(f, num=nx)
        cy = read_floats(f, num=ny)
        cz = read_floats(f, num=nz)

        ex = read_floats(f, num=nx)
        ey = read_floats(f, num=ny)
        ez = read_floats(f, num=nz)

        time = read_float(f)

        field_type = read_int(f)
        if field_type == 0:
            nvars = 1
        elif field_type == 1:
            nvars = ndims
        else:
            raise Exception(f"Unknown field_type value: {field_type}. Must be either 0 or 1")

        len_varnames = []
        for _ in range(nvars):
            len_varnames += [read_int(f)]

        varnames = []
        for nvar in range(nvars):
            varnames += [read_str(f, length=len_varnames[nvar])]

        values = np.fromfile(f, dtype=float_type, count=np.prod(dimsizes)).reshape(dimsizes + [1])
        for nvar in range(1,nvars):
            vadd = np.fromfile(f, dtype=float_type, count=np.prod(dimsizes)).reshape(dimsizes + [1])
            values = np.concatenate((values, vadd), axis=-1)

        f.close()

        I = nx - 2 * gcx
        J = ny - 2 * gcy
        K = nz - 2 * gcz

        self.data['cx'] = I
        self.data['cy'] = J
        self.data['cz'] = K

        self.variable_names = ['I', 'J', 'K']

        for i in range(nvars):
            self.variable_names.append(varnames[i][0])

        self.data['I'] = cx[gcx:nx - gcx]
        self.data['J'] = cy[gcy:ny - gcy]
        self.data['K'] = cz[gcz:nz - gcz]

        values = values.T

        for i in range(ndims, ndims + nvars):
            self.data[self.variable_names[i]] = values[i - ndims, gcz:nz - gcz, gcy:ny - gcy, gcx:nx - gcx]

        if ndims == 1:
            var_begin = 1
        elif ndims == 2:
            var_begin = 2
        elif ndims == 3:
            var_begin = 3

        if self.var == None:
            column_names = [self.variable_names[i] for i in range(var_begin, len(self.variable_names))] 
        else:
            column_names = self.var

        if self.mval != None:
            for col, abs_val in zip(column_names, self.mval):
                idx = np.where(np.abs(self.data[col]) > abs_val)
                self.data[col][idx] = np.nan

    def process_file(self, ndim, column_names, mval):
        filename_parts = os.path.splitext(self.filename)
        filename_extension = filename_parts[-1]
        self.var = column_names
        self.mval = mval

        # print(ndim, self.filename)

        if filename_extension == '.plt':
            self.__read_plt(ndim)
        elif filename_extension == '.nsx':
            self.__read_bin()

    def get_variable_names(self):
        filename_parts = os.path.splitext(self.filename)
        filename_extension = filename_parts[-1]

        if filename_extension == '.plt':
            self.__get_variable_names_plt()
        elif filename_extension == '.nsx':
            self.__get_variable_names_bin()