import h5py
import numpy as np
import re
from netCDF4 import Dataset
import glob
import argparse
import os
import warnings
from dataclasses import dataclass

attrs_dic = {
    "xt": {
        "long_name": "West-East displacement of cell centers",
        "units": "m",
    },
    "xm": {
        "long_name": "West-East displacement of cell edges",
        "units": "m",
    },
    "yt": {
        "long_name": "South-North displacement of cell centers",
        "units": "m",
    },
    "ym": {
        "long_name": "South-North displacement of cell edges",
        "units": "m",
    },
    "zm": {
        "long_name": "Vertical displacement of cell centers",
        "units": "m",
    },
    "zt": {
        "long_name": "Vertical displacement of cell edges",
        "units": "m",
    },
}


@dataclass
class settings:
    regex_pattern: re.Pattern
    all_x: bool = True
    all_y: bool = True
    all_z: bool = True
    has_z: bool = True
    has_z_soil: bool = False


pattern_dict = {
    "fielddump": settings(
        regex_pattern=re.compile(
            r".*/fielddump\.(?P<x_idx>\d+)\.(?P<y_idx>\d+)\.(?P<expnr>\d+)\.nc"
        )
    ),
    "cape": settings(
        regex_pattern=re.compile(
            r".*/cape\.x(?P<x_idx>\d+)y(?P<y_idx>\d+)\.(?P<expnr>\d+)\.nc"
        ),
        has_z=False,
    ),
    "surfcross": settings(
        regex_pattern=re.compile(
            r".*/surfcross\.x(?P<x_idx>\d+)y(?P<y_idx>\d+)\.(?P<expnr>\d+)\.nc"
        ),
        has_z=False,
    ),
    "slrbcross": settings(
        regex_pattern=re.compile(
            r".*/slrbcross\.x(?P<x_idx>\d+)y(?P<y_idx>\d+)\.(?P<expnr>\d+)\.nc"
        ),
        has_z=False,
        has_z_soil=True,
    ),
    "lsmcrossxy": settings(
        regex_pattern=re.compile(
            r".*/lsmcrossxy\.\d+\.x(?P<x_idx>\d+)y(?P<y_idx>\d+)\.(?P<expnr>\d+)\.nc"
        ),
        has_z=False,
        has_z_soil=False,
    ),
    "lsmcrossxz": settings(
        regex_pattern=re.compile(r".*/lsmcrossxz\.x(\d+)y(\d+)\.(\d+)\.nc"),
    ),
    "crossyz": settings(
        regex_pattern=re.compile(
            r".*/crossyz\.(?P<x_idx>\d+)\.y(?P<y_idx>\d+)\.(?P<expnr>\d+)\.nc"
        ),
        has_z=True,
        all_x=False,
    ),
    "crossxz": settings(
        regex_pattern=re.compile(
            r".*/crossxz\.(?P<y_idx>\d+)\.x(?P<x_idx>\d+)\.(?P<expnr>\d+)\.nc"
        ),
        has_z=True,
        all_y=False,
    ),
    "crossxy": settings(
        regex_pattern=re.compile(
            r".*/crossxy\.(?P<z_idx>\d+)\.x(?P<x_idx>\d+)y(?P<y_idx>\d+)\.(?P<expnr>\d+)\.nc"
        ),
        has_z=True,
        all_z=False,
    ),
}


def combine_fielddump(files, pattern_settings: settings, output_file):
    indices = []
    x_idx_ls = []
    y_idx_ls = []
    z_idx_ls = []
    expnr_idx_ls = []

    all_x = pattern_settings.all_x
    all_y = pattern_settings.all_y
    all_z = pattern_settings.all_z
    has_z = pattern_settings.has_z
    has_z_soil = pattern_settings.has_z_soil

    for f in files:
        match = pattern_settings.regex_pattern.match(f)
        x_idx_ls.append(int(match["x_idx"]))
        y_idx_ls.append(int(match["y_idx"]))

        if all_z:
            z_idx = 1
        else:
            z_idx = int(match["z_idx"])
            z_idx_ls.append(z_idx)
        indices.append(
            (int(match["x_idx"]), int(match["y_idx"]), int(match["expnr"]), z_idx)
        )

    x_idx_ls = sorted(list(set(x_idx_ls)))
    y_idx_ls = sorted(list(set(y_idx_ls)))
    z_idx_ls = sorted(list(set(z_idx_ls)))

    if all_x:
        nx = max(i for i, j, k, iz in indices) + 1
    else:
        nx = len(set(i for i, j, k, iz in indices))
    if all_y:
        ny = max(j for i, j, k, iz in indices) + 1
    else:
        ny = len(set(j for i, j, k, iz in indices))

    if all_z:
        nz = None
    else:
        nz = len(z_idx_ls)
    # --- inspect one file for metadata ---
    with Dataset(files[0]) as ds0:
        dim_order = list(ds0.dimensions.keys())
        varnames = list(ds0.variables.keys())
        dim_sizes = {d: len(ds0.dimensions[d]) for d in dim_order}
        dtypes = {v: ds0.variables[v].dtype for v in varnames}
        shapes = {v: ds0.variables[v].shape for v in varnames}
        var_dims = {v: ds0.variables[v].dimensions for v in varnames}
        vattrs = {v: dict(ds0.variables[v].__dict__) for v in varnames}
        gattrs = dict(ds0.__dict__)
        if has_z and all_z:
            zm_global = ds0.variables["zm"][:]
            zt_global = ds0.variables["zt"][:]
        if has_z_soil:
            zts_global = ds0.variables["zts"][:]
        time_global = ds0.variables["time"][:]

        if not all_z:
            dim_order.append("zt")
            varnames.append("zt")
            dim_sizes["zt"] = 1
            dtypes["zt"] = np.float32
            shapes["zt"] = (1,)
            var_dims["zt"] = ("zt",)
            vattrs["zt"] = attrs_dic["zt"]

            for var in var_dims:
                if var in ["xt", "yt", "xm", "ym", "zt", "zm", "time"]:
                    continue
                var_dims[var] = (*var_dims[var], "zt")
                shapes[var] = (*shapes[var], 1)
        if not all_x:
            dim_order.append("xt")
            dim_order.append("xm")
            varnames.append("xt")
            varnames.append("xm")
            dim_sizes["xt"] = 1
            dim_sizes["xm"] = 1
            dtypes["xt"] = np.float32
            dtypes["xm"] = np.float32
            shapes["xt"] = (1,)
            shapes["xm"] = (1,)
            var_dims["xt"] = ("xt",)
            var_dims["xm"] = ("xm",)
            vattrs["xt"] = attrs_dic["xt"]
            vattrs["xm"] = attrs_dic["xm"]

            for var in var_dims:
                if var in ["xt", "yt", "xm", "ym", "zt", "zm", "time"]:
                    continue
                var_dims[var] = (*var_dims[var], "xt")
                shapes[var] = (*shapes[var], 1)
        if not all_y:
            dim_order.append("yt")
            dim_order.append("ym")
            varnames.append("yt")
            varnames.append("ym")
            dim_sizes["yt"] = 1
            dim_sizes["ym"] = 1
            dtypes["yt"] = np.float32
            dtypes["ym"] = np.float32
            shapes["yt"] = (1,)
            shapes["ym"] = (1,)
            var_dims["yt"] = ("yt",)
            var_dims["ym"] = ("ym",)
            vattrs["yt"] = attrs_dic["yt"]
            vattrs["ym"] = attrs_dic["ym"]
            for var in var_dims:
                if var in ["xt", "yt", "xm", "ym", "zt", "zm", "time"]:
                    continue
                var_dims[var] = (*var_dims[var], "yt")
                shapes[var] = (*shapes[var], 1)

    # tiled dims
    x_dim, y_dim = "xt", "ym"
    x_tile, y_tile = dim_sizes[x_dim], dim_sizes[y_dim]
    global_x, global_y = nx * x_tile, ny * y_tile

    if not all_z:
        global_z = nz

    # --- stitch global coordinate arrays ---
    # assume each tile has 1D coords "xt", "ym" valid for its local chunk
    xt_global = np.full(global_x, np.nan)
    xm_global = np.full(global_x, np.nan)
    yt_global = np.full(global_y, np.nan)
    ym_global = np.full(global_y, np.nan)
    if not all_z:
        zt_global = np.full(global_z, np.nan)
        zm_global = np.full(global_z, np.nan)

    for f, (ix, iy, ik, iz) in zip(files, indices):
        with Dataset(f) as ds:
            xt_local = ds.variables["xt"][:]
            xm_local = ds.variables["xm"][:]
            yt_local = ds.variables["yt"][:]
            ym_local = ds.variables["ym"][:]
            if not all_z:
                try:
                    zt_local = ds.variables["zt"][:]
                    zm_local = ds.variables["zm"][:]
                except KeyError:
                    zt_local = np.array([iz])
                    zm_local = np.array([iz])
        if all_x:
            x0 = ix * x_tile
        else:
            x0 = x_idx_ls.index(ix)
        if all_y:
            y0 = iy * y_tile
        else:
            y0 = y_idx_ls.index(iy)
        if not all_z:
            z0 = z_idx_ls.index(iz)
        xt_global[x0 : x0 + x_tile] = xt_local
        xm_global[x0 : x0 + x_tile] = xm_local
        ym_global[y0 : y0 + y_tile] = ym_local
        yt_global[y0 : y0 + y_tile] = yt_local
        if not all_z:
            zt_global[z0 : z0 + 1] = zt_local
            zm_global[z0 : z0 + 1] = zm_local

    # --- create stitched VDS file ---
    with h5py.File(output_file, "w", libver="latest") as fout:
        for k, val in gattrs.items():
            if not (type(val) is str):
                fout.attrs[k] = val
            else:
                fout.attrs[k] = np.array(
                    val, dtype=h5py.string_dtype("ascii", len(str(val)))
                )

        # coordinate variables with real values
        fout.create_dataset("xt", data=xt_global).make_scale("xt")
        if has_z and all_z:
            fout.create_dataset("zt", data=zt_global).make_scale("zt")
            fout.create_dataset("zm", data=zm_global).make_scale("zm")
        if has_z_soil:
            fout.create_dataset("zts", data=zts_global).make_scale("zts")

        if not all_z:
            fout.create_dataset("zt", data=zt_global).make_scale("zt")
            fout.create_dataset("zm", data=zm_global).make_scale("zm")
        fout.create_dataset("xm", data=xm_global).make_scale("xm")
        fout.create_dataset("yt", data=yt_global).make_scale("yt")
        fout.create_dataset("ym", data=ym_global).make_scale("ym")
        fout.create_dataset("time", data=time_global).make_scale("time")

        # any other dims copied as indices
        for d in dim_order:
            if d in ["xt", "yt", "xm", "ym", "zm", "zt", "zts", "time"]:
                continue
            fout.create_dataset(d, data=np.arange(dim_sizes[d])).make_scale(d)

        # stitch variables
        for v in varnames:
            if not (
                v
                in [
                    "xslice_t",
                    "yslice_t",
                    "xslice_m",
                    "yslice_m",
                    "xt",
                    "yt",
                    "xm",
                    "ym",
                    "zm",
                    "zt",
                    "zts",
                    "time",
                ]
            ):
                # skip coords already created
                if v in ["zslice_t", "zslice_m"]:
                    continue
                shape = shapes[v]
                dims = var_dims[v]
                if any(dim in dims for dim in ["xt", "yt", "xm", "ym"]):
                    for dim in ["xt", "yt", "xm", "ym"]:
                        if dim in dims:
                            if "x" in dim:
                                var_x_dim = dim
                            else:
                                var_y_dim = dim
                        if not all_x:
                            var_x_dim = "xt"
                        if not all_y:
                            var_y_dim = "yt"

                    # find axis positions
                    x_axis = dims.index(var_x_dim)
                    y_axis = dims.index(var_y_dim)
                    if not all_z:
                        z_axis = dims.index("zt")

                    # replace tiled dims with global sizes
                    global_shape = list(shape)
                    global_shape[x_axis] = global_x
                    global_shape[y_axis] = global_y
                    if not all_z:
                        global_shape[z_axis] = global_z
                    global_shape = tuple(global_shape)

                    layout = h5py.VirtualLayout(shape=global_shape, dtype=dtypes[v])

                    for f, (ix, iy, ik, iz) in zip(files, indices):
                        vsource = h5py.VirtualSource(f, v, shape=shape)
                        if all_x:
                            x0 = ix * x_tile
                        else:
                            x0 = x_idx_ls.index(ix)
                        if all_y:
                            y0 = iy * y_tile
                        else:
                            y0 = y_idx_ls.index(iy)
                        if not all_z:
                            z0 = z_idx_ls.index(iz)

                        # slice tuple dynamically
                        slices = [slice(None)] * len(shape)
                        slices[x_axis] = slice(x0, x0 + x_tile)
                        slices[y_axis] = slice(y0, y0 + y_tile)
                        if not all_z:
                            slices[z_axis] = slice(z0, z0 + 1)
                        layout[tuple(slices)] = vsource

                    dset = fout.create_virtual_dataset(v, layout, fillvalue=np.nan)

                else:
                    # not tiled -> copy from one file
                    with Dataset(files[0]) as ds0:
                        data = ds0.variables[v][:]
                    dset = fout.create_dataset(v, data=data)

                # attach dimension scales
                for i, d in enumerate(var_dims[v]):
                    dset.dims[i].attach_scale(fout[d])

            # copy variable attributes
            for k, val in vattrs[v].items():
                if not (type(val) is str):
                    fout[v].attrs[k] = val
                else:
                    fout[v].attrs[k] = np.array(
                        val, dtype=h5py.string_dtype("ascii", len(str(val)))
                    )


def determine_output_types(in_path):
    glob_patterns = []
    files = glob.glob(f"{in_path}/*.nc")
    recognized_types = [
        "cape",
        "fielddump",
        "surfcross",
        "slrbcross",
        "crossyz",
        "crossxz",
        "crossxy",
    ]
    not_recognised_types = ["lsmcrossxz", "lsmcrossxy"]
    for output_type in recognized_types:
        if any(
            (
                (
                    (output_type in file)
                    and not (
                        any(not_recog in file for not_recog in not_recognised_types)
                    )
                )
                for file in files
            )
        ):
            pattern = f"{in_path}/{output_type}*.nc"
            if (output_type, pattern, f"{output_type}.nc") in glob_patterns:
                continue
            glob_patterns.append((output_type, pattern, f"{output_type}.nc"))
    return glob_patterns


def combine_run(in_path: str) -> None:
    """Combine all recognized DALES files in a run folder.

    This mirrors the script behavior so it can be called from other
    Python code, e.g. ``import dales_tile_combiner; dales_tile_combiner.combine_run("run_001")``.
    """

    glob_patterns = determine_output_types(in_path)

    for output_type, pattern, out_file in glob_patterns:
        files = list(glob.glob(pattern))
        pattern_settings = pattern_dict[output_type]
        combine_fielddump(files, pattern_settings, out_file)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Combine MPI-split fielddump into one Virtual Dataset file."
    )
    parser.add_argument(
        "in_path", help="Path containing input files (e.g. 'run_001/*.nc')"
    )
    args = parser.parse_args()
    combine_run(args.in_path)


if __name__ == "__main__":
    main()
