import argparse
import multiprocessing as mp
import os
import pickle
import re
import warnings
from subprocess import call
from time import time

import OpenPNM
import numpy as np
from glob import glob

import sys

sys.path.append('../')

from textures3d.visualization import load_3d_hdf5
from textures3d.kern_utils import postprocess


def clear_folder(samplename, filepath="./network_extraction/"):
    file_pattern = r'{}_'.format(samplename)
    for File in os.listdir(filepath):
        if re.search(file_pattern, File):
            os.remove("{}{}".format(filepath, File))


def splits_to_files(split, samplename, filepath="./network_extraction/"):
    # file_pattern = r'{}_\d+'.format(samplename)
    # flag = True
    clear_folder(samplename, filepath)
    for i in range(split.shape[0]):
        cube = split[i, :, :, :]
        vector = cube.flatten()
        vector.tofile("{}{}_{}.raw".format(filepath, samplename, i))


def compute_network(filename, prec, minr2, filepath="./network_extraction/", cubesize=100,
                    config_name="input.dat", output_path="./networks/"):
    input_template = pickle.load(open("input_template.p", mode="rb"))
    fullpath = "{}{}".format(filepath, filename)
    samplename = filename[:-4]
    input_template = input_template.format(fullpath, cubesize, cubesize, cubesize,
                                           prec, minr2)

    try:
        text_file = open("{}{}".format(filepath, config_name), "w")
        text_file.write(input_template)
        text_file.close()
    except IOError:
        print("Error, file could not be created")

    call(["{}ppd".format(filepath), "{}{}".format(filepath, config_name)])
    call(["{}porenet".format(filepath), "{}{}".format(filepath, config_name)])

    output_patterns = ["{}.orig.dat", "{}.ppd", "{}_link1.dat", "{}_link2.dat",
                       "{}_node1.dat", "{}_node2.dat"]
    for pattern in output_patterns:
        os.rename("{}{}".format(filepath, pattern.format(samplename)),
                  "{}{}".format(output_path, pattern.format(samplename)))


def compute_networks(samplename, filepath="./network_extraction/", cubesize=100,
                     config_name="input.dat"):
    clear_folder(samplename, filepath="./networks/")
    pattern = r'{}_[0-9]+.raw'.format(samplename)
    for File in os.listdir(filepath):
        if re.search(pattern, File):
            compute_network(File, 5.345, 2, filepath, cubesize, config_name)


def load_network(prefix, filepath="./networks/"):
    """Loads network file of Statoil format, and trims isolated pores
       Returns OpenPNM Network class"""
    net = OpenPNM.Utilities.IO.Statoil.load(filepath, prefix)
    # Removing isolated pores
    pores_to_ignore = net.check_network_health()["trim_pores"]
    OpenPNM.Network.tools.trim(net, pores=pores_to_ignore)
    return net


def compute_permeabilities(net, input_pores, output_pores, T=298.0, viscosity=0.001, mu_w=0.001,
                           input_pressure=101325, out_pressure=202650):
    """Computes permeability in mD
       input_pores, output_pores are boolean lists, used for indexing
       net is a Statoil format network with already defined pore.radius, pore.volume,
       throat.radius, throat.length, throat.volume arrays"""
    mD_per_m2 = 1.013249966e+12 * 1000
    # Creating and updating corresponding geometry class
    geom = OpenPNM.Geometry.GenericGeometry(network=net, pores=net.Ps, throats=net.Ts)
    geom['pore.diameter'] = net["pore.radius"] * 2
    geom['pore.volume'] = net["pore.volume"]
    geom['throat.diameter'] = net["throat.radius"] * 2
    geom['throat.length'] = net["throat.length"]
    geom['throat.volume'] = net["throat.volume"]
    # Creating water phase
    water = OpenPNM.Phases.GenericPhase(network=net)
    water['pore.temperature'] = T
    water['pore.viscosity'] = viscosity
    # Creating physics class
    phys_water = OpenPNM.Physics.GenericPhysics(network=net, phase=water, geometry=geom)
    R = geom['throat.diameter'] / 2
    L = geom['throat.length']
    phys_water['throat.hydraulic_conductance'] = 3.14159 * R ** 4 / (8 * mu_w * L)
    # Creating algorithm class
    alg = OpenPNM.Algorithms.StokesFlow(network=net, phase=water)
    BC1_pores = input_pores
    alg.set_boundary_conditions(bctype='Dirichlet', bcvalue=input_pressure,
                                pores=BC1_pores)
    BC2_pores = output_pores
    alg.set_boundary_conditions(bctype='Dirichlet', bcvalue=out_pressure,
                                pores=BC2_pores)
    try:
        alg.run()
        k = alg.calc_eff_permeability()  # * mD_per_m2
    except Exception as e:
        print(e)
        k = np.nan
    return k


def water_permeability(prefix, filepath="./networks/"):
    """Wrapper method for computing permeability with default settings
       Computes permeability for input and output pores, specified in Statoil format,
       for water of temperature 298K and viscosity 0.001 with pressure difference between
       input and output of 101325 pascal.
       Returns permeability in mD"""
    try:
        net = load_network(prefix, filepath)
    except Exception as e:
        print(e)
        return np.nan
    input_pores = net.pores('inlets')
    output_pores = net.pores('outlets')
    return compute_permeabilities(net, input_pores, output_pores)


def water_permeabilites(samplename, split_num, filepath="./networks/", to_file=True,
                        output="./results/", supress_warnings=True):
    """Computes permeabilities for all prefixes of type samplename_i, where i is an integer number
       If to_file is true, writes output to file
       Returns array of permeabilities"""
    warnings.filterwarnings('ignore')
    permeabilities = np.zeros(split_num)
    for i in range(split_num):
        permeabilities[i] = water_permeability("{}_{}".format(samplename, i), filepath)
        print(permeabilities)
    if to_file:
        np.save("{}{}_permeabilities".format(output, samplename), permeabilities)
    return permeabilities


def write_metadata(samplename, cubesize, shift, nm_per_voxel, split_num, output="./results/"):
    metadata_template = pickle.load(open("metadata_template.p", mode="rb"))
    metadata = metadata_template.format(samplename, cubesize, shift, nm_per_voxel, split_num)
    f = open("{}{}_data.txt".format(output, samplename), "w")
    f.write(metadata)
    f.close()


def compute_features(cube, dict):
    """Minkowski feature computation with precomputed updates
       Returns V, S, B, X features"""
    cube = np.pad(cube, 1, "constant")
    n_0, n_1, n_2, n_3 = 0, 0, 0, 0
    for x in range(1, cube.shape[0] - 1):
        for y in range(1, cube.shape[1] - 1):
            for z in range(1, cube.shape[-1] - 1):
                dn_3, dn_2, dn_1, dn_0 = dict[
                    hash(tuple(cube[x - 1:x + 2, y - 1:y + 2, z - 1:z + 1].flatten()))]
                n_3 += dn_3
                n_2 += dn_2
                n_1 += dn_1
                n_0 += dn_0
    V = n_3
    S = -6 * n_3 + 2 * n_2
    B = 3 * n_3 / 2 - n_2 + n_1 / 2
    X = - n_3 + n_2 - n_1 + n_0
    return V, S, B, X


def compute_batch(batch):
    path = "../slice2pores/statistics/comb_to_update_new.p"
    dict = pickle.load(open(path, "rb"))
    features = list(map(lambda b: compute_features(b, dict), batch))
    return features


def compute_batch_parallel(batch, num_threads=4):
    slice_size = round(len(batch) / num_threads + 0.5)
    starts = list(range(0, num_threads * slice_size, slice_size))
    batches = [batch[start:start + slice_size] for start in starts]
    pool = mp.Pool(num_threads)
    result = []
    res = [pool.apply_async(compute_batch, args=(b,), callback=lambda lst: result.extend(lst)) for b in batches]
    pool.close()
    pool.join()
    return result


def get_permeability(arr, size=128):
    res = []
    for j in range(len(arr)):
        print('kern %s' % j)
        split = arr[j].reshape(1, size, size, size)
        split_num = split.shape[0]
        split = np.uint8(split)
        splits_to_files(split, 'samplename')
        compute_networks("samplename", cubesize=size)
        p = water_permeabilites('samplename', split_num, to_file=False)
        res.append(p)

        print(res)

    res = np.array([r[0] for r in res])

    return res


parser = argparse.ArgumentParser()
parser.add_argument('--dir', required=True, type=str, help='folder hdf5 from exman experiment (with files gen_*.hdf5, orig_*.hdf5')
parser.add_argument('--sample_name', required=True, type=str, help='I.e. "Berea"')
parser.add_argument('--size', required=True, type=int, help='Size of samples')
parser.add_argument('--inv', action='store_true', dest='inv', default=False, help='Size of samples')


def main():
    opt = parser.parse_args()

    ORIGINAL = False

    if not os.path.exists('./networks'):
        os.makedirs('./networks')

    # original_f = glob(os.path.join(opt.dir, 'orig*.hdf5'))
    gen_f = glob(os.path.join(opt.dir, 'gen*.hdf5'))[:100]

    # print('Original samples %s, gen samples %s' % (len(original_f), len(gen_f)))

    # original = np.vstack([load_3d_hdf5(f) for f in original_f])
    gen = np.vstack([load_3d_hdf5(f) for f in gen_f])

    # original = original.reshape(len(original_f), 1, opt.size, opt.size, opt.size)
    gen = gen.reshape(len(gen_f), 1, opt.size, opt.size, opt.size)

    # original = original / 255
    gen = postprocess(gen)
    # print(original.max(), original.min())
    print(gen.max(), gen.min())

    if opt.inv:
        print('Inverting %s' % opt.sample_name)
        gen = 1 - gen


    gen_permeability = get_permeability(gen, size=opt.size)
    np.savetxt('./res_200/baseline_%s_gen.txt' % opt.sample_name, gen_permeability)


if __name__ == '__main__':
    main()