import json
import os
import pathlib
from itertools import groupby
from operator import attrgetter
from typing import TYPE_CHECKING

import numpy as np

from ada import Beam, Material, Part, Section
from ada.concepts.containers import Nodes
from ada.config import get_logger

from .utils import to_safe_name

if TYPE_CHECKING:
    from ada import FEM, Assembly

logger = get_logger()


def write_assembly_to_cache(assembly: "Assembly", cache_file_path):
    """Write the Assembly information to a HDF5 file format for High performance cache."""
    import h5py

    # TODO: Add support for FEM only to begin with

    cache_file_path = pathlib.Path(cache_file_path)
    os.makedirs(cache_file_path.parent, exist_ok=True)
    h5_filename = cache_file_path.with_suffix(".h5")
    with h5py.File(h5_filename, "w") as f:
        info = f.create_group("INFO")
        info.attrs.create("NAME", assembly.name)

        parts_group = f.create_group("PARTS")

        walk_parts(parts_group, assembly)

    print(f'Saved cached model at "{cache_file_path}"')


def walk_parts(cache_p, part):
    for p in part.parts.values():
        part_group = add_part_to_cache(p, cache_p)
        part_group.attrs.create("PARENT", to_safe_name(p.parent.name))
        walk_parts(part_group, p)


def add_part_to_cache(part: Part, parent_part_group):
    part_group = parent_part_group.create_group(to_safe_name(part.name))

    part_group.attrs.create("METADATA", json.dumps(part.metadata))

    if len(part.nodes) > 0:
        add_nodes_to_cache(part.nodes, part_group)

    if len(part.sections) > 0:
        add_sections_to_cache(part, part_group)

    if len(part.materials) > 0:
        add_materials_to_cache(part, part_group)

    if len(part.beams) > 0:
        add_beams_to_cache(part, part_group)

    if len(part.plates) > 0:
        add_plates_to_cache()

    if len(part.shapes) > 0:
        add_shapes_to_cache()

    if len(part.pipes) > 0:
        add_pipes_to_cache()

    if len(part.walls) > 0:
        add_walls_to_cache()

    # Add FEM object
    if len(part.fem.nodes) > 0:
        print(f'Caching FEM data from "{part.name}"')
        add_fem_to_cache(part.fem, part_group)

    return part_group


def add_sections_to_cache(part, parts_group):
    prefix = "SECTIONS"

    def add_ints_to_cache(s: Section):
        return [x if x is not None else 0 for x in [s.r, s.wt, s.h, s.w_top, s.w_btn, s.t_w, s.t_ftop, s.t_fbtn, s.id]]

    def add_strings_to_cache(s: Section):
        return [s.guid, s.name, s.units.value, s.type.value]

    parts_group.create_dataset(f"{prefix}_STR", data=[add_strings_to_cache(bm) for bm in part.sections])
    parts_group.create_dataset(f"{prefix}_INT", data=[add_ints_to_cache(bm) for bm in part.sections])


def add_materials_to_cache(part, parts_group):
    prefix = "MATERIALS"

    def add_ints_to_cache(e: Material):
        m = e.model
        return [m.E, m.rho, m.sig_y, e.id]

    def add_strings_to_cache(e: Material):
        return [e.guid, e.name, e.units.value]

    parts_group.create_dataset(f"{prefix}_INT", data=[add_ints_to_cache(bm) for bm in part.materials])
    parts_group.create_dataset(f"{prefix}_STR", data=[add_strings_to_cache(bm) for bm in part.materials])


def add_plates_to_cache():
    logger.error("Plate caching is not yet implemented")


def add_shapes_to_cache():
    logger.error("Shape caching is not yet implemented")


def add_pipes_to_cache():
    logger.error("Pipes caching is not yet implemented")


def add_walls_to_cache():
    logger.error("Walls caching is not yet implemented")


def add_beams_to_cache(part: Part, parts_group):
    prefix = "BEAMS"

    def add_int_cache(bm: Beam, up=False):
        if up is True:
            nids = bm.up.tolist()
        else:
            nids = [bm.n1.id, bm.n2.id]
        if None in nids:
            raise ValueError()
        return nids

    def add_str_cache(bm: Beam):
        return [bm.guid, bm.name, bm.section.name, bm.material.name, json.dumps(bm.metadata)]

    parts_group.create_dataset(f"{prefix}_INT", data=[add_int_cache(bm) for bm in part.beams])
    parts_group.create_dataset(f"{prefix}_STR", data=[add_str_cache(bm) for bm in part.beams])
    parts_group.create_dataset(f"{prefix}_UP", data=[add_int_cache(bm, True) for bm in part.beams])


def add_fem_to_cache(fem: "FEM", part_group):
    fem_group = part_group.create_group("FEM")
    fem_group.attrs.create("NAME", to_safe_name(fem.name))

    # Add Nodes
    add_nodes_to_cache(fem.nodes, fem_group)

    # Add elements
    elements_group = fem_group.create_group("MESH")
    for group, elements in groupby(sorted(fem.elements, key=attrgetter("type")), key=attrgetter("type")):
        med_cells = elements_group.create_group(group.value)
        med_cells.create_dataset("ELEMENTS", data=[[int(el.id), *[int(n.id) for n in el.nodes]] for el in elements])


def add_nodes_to_cache(nodes: Nodes, group):
    points = np.array([[n.id, *n.p] for n in nodes])
    coo = group.create_dataset("NODES", data=points)
    coo.attrs.create("NBR", len(points))
