import os
import numpy as np
from datetime import datetime


def zone_to_zonn(zonefile):
    zonnfile = os.path.splitext(zonefile)[0] + ".zonn"
    with open(zonefile, "r") as fh:
        fh.readline()
        lns = fh.readlines()
    with open(zonnfile, "w") as fout:
        fout.write("zonn\n")
        for l in lns:
            fout.write(l)


def spherical_writeFEHM(node_locations, filename_base, title="default"):
    """
    Create FEHM .stor and .fehmn input files for 1d spherical simulations.
    Assumes logical 1D structure.
    :arg node_locations: radial location of nodes
    :type node_locations: array(float)
    :arg filename_base: base name for the .stor and .fehmn files
    :type filename_base: str
    :arg title: optional title for simulation
    :type title: str

    Example:
    from pylagrit import utilities as util
    import numpy as np
    nodes = np.linspace(0.0,200.0)
    filename_base = "test"
    title = "test sim"
    # write .stor and .fehmn files
    util.spherical_writeFEHM(nodes,filename_base,title)
    """
    # stor file header
    now = datetime.now()
    print_datetime = now.strftime("%m/%d/%Y  %H:%M:%S")
    header = (
        "FEHM .stor file generated by PyLaGriT  "
        + print_datetime
        + "\n"
        + "title:  "
        + title
        + "\n"
    )

    # calculate dx, face locations, face areas, volumes, and geometric coefficients
    dx = spherical_dx(node_locations)
    faces = spherical_faces(node_locations)
    area = spherical_areas(faces)
    volumes = spherical_volumes(node_locations)
    coeffs = area / dx

    # matrix metadata header info
    neq = np.size(node_locations)
    num_coeffs = neq - 1
    FEHM_mem = (neq - 2) * 3 + 4 + neq + 1
    num_area_coeff = 1
    max_connect = 3
    matrix_params = (num_coeffs, neq, FEHM_mem, num_area_coeff, max_connect)
    params_header = "        " + "        ".join(str(s) for s in matrix_params) + "\n"

    # row count
    row_count = np.empty([neq + 1])
    start = neq + 1
    row_count[0] = start
    row_count[1] = row_count[0] + 2
    for i in range(2, neq):
        row_count[i] = row_count[i - 1] + 3
    row_count[neq] = row_count[neq - 1] + 2

    # row entries
    row_entries = []
    row_entries.append(1)
    row_entries.append(2)
    for i in range(1, neq - 1):
        row_entries.append(i)
        row_entries.append(i + 1)
        row_entries.append(i + 2)
    row_entries.append(neq - 1)
    row_entries.append(neq)

    # index into geometric coefficient matrix
    coeff_indices = []
    coeff_indices.append(0)
    coeff_indices.append(1)
    for i in range(1, neq - 1):
        coeff_indices.append(i)
        coeff_indices.append(0)
        coeff_indices.append(i + 1)
    coeff_indices.append(neq - 1)
    coeff_indices.append(0)
    # neq + 1 extra 0s for padding
    for i in range(0, neq + 1):
        coeff_indices.append(0)

    # index of diagonal entries
    diagonal_indices = np.empty([neq])
    diagonal_indices[0] = neq + 2
    for i in range(1, neq):
        diagonal_indices[i] = diagonal_indices[i - 1] + 3

    # open stor file, write header and matrix parameters
    storfile = filename_base + ".stor"
    sfile = open(storfile, "w")
    _ = sfile.write(header)
    _ = sfile.write(params_header)

    # write Voronoi volumes
    count = 1
    for vol in volumes:
        print("  %1.12e" % vol, end="" if count % 5 else "\n", file=sfile)
        count += 1
    if (count - 1) % 5:
        _ = sfile.write("\n")

    # write row counts
    count = 1
    for row in row_count:
        print(" %9d" % row, end="" if count % 5 else "\n", file=sfile)
        count += 1
    if (count - 1) % 5:
        _ = sfile.write("\n")

    # write row entries
    count = 1
    for row in row_entries:
        print(" %9d" % row, end="" if count % 5 else "\n", file=sfile)
        count += 1
    if (count - 1) % 5:
        _ = sfile.write("\n")

    # write geometric coefficient indices
    count = 1
    for idx in coeff_indices:
        print(" %9d" % idx, end="" if count % 5 else "\n", file=sfile)
        count += 1
    if (count - 1) % 5:
        _ = sfile.write("\n")

    # write diagonal indices
    count = 1
    for idx in diagonal_indices:
        print(" %9d" % idx, end="" if count % 5 else "\n", file=sfile)
        count += 1
    if (count - 1) % 5:
        _ = sfile.write("\n")

    # write geometric coefficients
    count = 1
    for coef in coeffs:
        print("  %1.12e" % coef, end="" if count % 5 else "\n", file=sfile)
        count += 1

    # close stor file
    if (count - 1) % 5:
        _ = sfile.write("\n")
    _ = sfile.close()

    # open fehmn file
    fehmnfile = filename_base + ".fehmn"
    ifile = open(fehmnfile, "w")
    # write header and neq
    _ = ifile.write("coor\n%d\n" % neq)
    # write node number and radial location
    for i in range(1, neq + 1):
        print(
            "        %3d        %12f        0        0" % (i, node_locations[i - 1]),
            file=ifile,
        )
    # write connectivity
    _ = ifile.write("\nelem\n")
    _ = ifile.write("%d %d\n" % (2, neq - 1))
    for i in range(1, neq):
        print("%3d   %3d   %3d" % (i, i, i + 1), file=ifile)
    # close fehmn file
    _ = ifile.write("\nstop\n")
    _ = ifile.close()


def spherical_faces(node_locations):
    """
    Calculate radial interface locations given radial node locations.
    Assumes 1st and last nodes lie on domain boundary.
    :arg node_locations: radial location of nodes
    :type node_locations: array_like(float)
    Returns: array of radial interface locations of size node_locations - 1
    """
    interface_locations = node_locations[:-1] + np.diff(node_locations) / 2.0
    return interface_locations


def spherical_areas(interface_locations):
    """
    Calculate spherical area of each interface given radial interface locations.
    :arg interface_locations: radial location of interfaces
    :type interface_locations: array_like(float)
    Returns: array of spherical interface areas of size interface_locations
    """
    areas = 4.0 * np.pi * interface_locations ** 2
    return areas


def spherical_dx(node_locations):
    """
    Calculate Delaunay edge lengths given radial node locations.
    :arg node_locations: radial location of nodes
    :type node_locations: array_like(float)
    Returns: array of Delaunay edge lengths of size node_locations - 1
    """
    delaunay = np.diff(node_locations)
    return delaunay


def spherical_volumes(node_locations):
    """
    Calculate Voronoi volume associated with each node given radial node locations.
    :arg node_locations: radial location of nodes
    :type node_locations: array_like(float)
    Returns: array of Voronoi volumes of size node_locations
    """
    coeff = np.pi * 4.0 / 3.0
    edges = spherical_faces(node_locations)
    edges = np.insert(edges, 0, node_locations[0])
    edges = np.append(edges, node_locations[np.size(node_locations) - 1])
    spheres = coeff * edges ** 3
    volumes = np.diff(spheres)
    assert np.all(volumes > 0), "ERROR: Negative volumes are not good."
    return volumes