#!/usr/bin/env python3
import argparse
import gc
import typing as t
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from astropy import units as u

import artistools as at


def make_cone(args):
    print("Making cone")

    angle_of_cone = 30  # in deg

    theta = np.radians([angle_of_cone / 2])  # angle between line of sight and edge is half angle of cone

    # merge_dfs, args.t_model, args.vmax = at.inputmodel.get_modeldata_tuple(args.modelpath[0], dimensions=3, get_elemabundances=True)
    merge_dfs = at.inputmodel.get_3d_model_data_merged_model_and_abundances_minimal(args)

    if args.positive_axis:
        cone = merge_dfs.loc[
            merge_dfs[f"cellpos_in[{args.sliceaxis}]"]
            >= 1
            / (np.tan(theta))
            * np.sqrt(
                (merge_dfs[f"cellpos_in[{args.other_axis2}]"]) ** 2
                + (merge_dfs[f"cellpos_in[{args.other_axis1}]"]) ** 2
            )
        ]  # positive axis
    else:
        cone = merge_dfs.loc[
            merge_dfs[f"cellpos_in[{args.sliceaxis}]"]
            <= -1
            / (np.tan(theta))
            * np.sqrt(
                (merge_dfs[f"cellpos_in[{args.other_axis2}]"]) ** 2
                + (merge_dfs[f"cellpos_in[{args.other_axis1}]"]) ** 2
            )
        ]  # negative axis
    # print(cone.loc[:, :[f'pos_{slice_on_axis}']])

    del merge_dfs  # merge_dfs not needed anymore so free memory
    gc.collect()

    return cone


def get_profile_along_axis(args=None, modeldata=None, derived_cols=False):
    print("Getting profile along axis")

    # merge_dfs, args.t_model, args.vmax = at.inputmodel.get_modeldata_tuple(args.modelpath, dimensions=3, get_elemabundances=True)
    if modeldata is None:
        modeldata, _ = at.inputmodel.get_modeldata(args.modelpath, get_elemabundances=True, derived_cols=derived_cols)

    position_closest_to_axis = modeldata.iloc[(modeldata[f"pos_{args.other_axis2}_min"]).abs().argsort()][:1][
        f"pos_{args.other_axis2}_min"
    ].item()

    if args.positive_axis:
        profile1D = modeldata.loc[
            (modeldata[f"pos_{args.other_axis1}_min"] == position_closest_to_axis)
            & (modeldata[f"pos_{args.other_axis2}_min"] == position_closest_to_axis)
            & (modeldata[f"pos_{args.sliceaxis}_min"] > 0)
        ]
    else:
        profile1D = modeldata.loc[
            (modeldata[f"pos_{args.other_axis1}_min"] == position_closest_to_axis)
            & (modeldata[f"pos_{args.other_axis2}_min"] == position_closest_to_axis)
            & (modeldata[f"pos_{args.sliceaxis}_min"] < 0)
        ]

    return profile1D.reset_index(drop=True)


def make_1D_profile(args):
    if args.makefromcone:
        cone = make_cone(args)

        slice1D = cone.groupby([f"cellpos_in[{args.sliceaxis}]"], as_index=False).mean()
        # where more than 1 X value, average rows eg. (1,0,0) (1,1,0) (1,1,1)

    else:  # make from along chosen axis
        slice1D = get_profile_along_axis(args)

    slice1D[f"cellpos_in[{args.sliceaxis}]"] = slice1D[f"cellpos_in[{args.sliceaxis}]"].apply(
        lambda x: x / args.t_model * (u.cm / u.day).to("km/s")
    )  # Convert positions to velocities
    slice1D = slice1D.rename(columns={f"cellpos_in[{args.sliceaxis}]": "vout_kmps"})
    # Convert position to velocity

    slice1D = slice1D.drop(
        ["inputcellid", f"cellpos_in[{args.other_axis1}]", f"cellpos_in[{args.other_axis2}]"], axis=1
    )  # Remove columns we don't need

    slice1D["rho_model"] = slice1D["rho_model"].apply(lambda x: np.log10(x) if x != 0 else -100)
    # slice1D = slice1D[slice1D['rho_model'] != -100]  # Remove empty cells
    # TODO: fix this, -100 probably breaks things if it's not one of the outer cells that gets chopped
    slice1D = slice1D.rename(columns={"rho_model": "log_rho"})

    slice1D.index += 1

    if not args.positive_axis:
        # Invert rows and *velocity by -1 to make velocities positive for slice on negative axis
        slice1D.iloc[:] = slice1D.iloc[::-1].to_numpy()
        slice1D["vout_kmps"] = slice1D["vout_kmps"].apply(lambda x: x * -1)

    # with pd.option_context('display.max_rows', None, 'display.max_columns', None):
    #     print(slice1D)

    # print(slice1D.keys())
    return slice1D


def make_1D_model_files(args):
    slice1D = make_1D_profile(args)
    query_abundances_positions = slice1D.columns.str.startswith("X_")
    model_df = slice1D.loc[:, ~query_abundances_positions]
    abundances_df = slice1D.loc[:, query_abundances_positions]

    # with pd.option_context('display.max_rows', None, 'display.max_columns', None):  # Print all rows in df
    #     print(model_df)

    # print(modelpath)
    model_df = model_df.round(decimals=5)  # model files seem to be to 5 sf
    model_df.to_csv(args.modelpath[0] / "model_1D.txt", sep=" ", header=False)  # write model.txt

    abundances_df.to_csv(args.modelpath[0] / "abundances_1D.txt", sep=" ", header=False)  # write abundances.txt

    with Path(args.modelpath[0], "model_1D.txt").open("r+") as f:  # add number of cells and tmodel to start of file
        content = f.read()
        f.seek(0, 0)
        f.write(f"{model_df.shape[0]}\n{args.t_model}".rstrip("\r\n") + "\n" + content)

    print("Saved abundances_1D.txt and model_1D.txt")


# with open(args.modelpath[0]/"model

# print(cone)

# cone = (merge_dfs.loc[merge_dfs[f'pos_{args.other_axis2}'] <= - (1/(np.tan(theta))
# * np.sqrt((merge_dfs[f'pos_{slice_on_axis}'])**2 + (merge_dfs[f'pos_{args.other_axis1}'])**2))])
# cone = merge_dfs
# cone = cone.loc[cone['rho_model'] > 0.0]


def make_plot(args):
    cone = make_cone(args)

    cone = cone.loc[cone["rho_model"] > 0.0002]  # cut low densities (empty cells?) from plot
    fig = plt.figure()
    ax = fig.gca(projection="3d")

    # print(cone['rho_model'])

    # set up for big model. For scaled down artis input model switch x and z
    x = cone["pos_z_min"].apply(lambda x: x / args.t_model * (u.cm / u.day).to("km/s")) / 1e3
    y = cone["pos_y_min"].apply(lambda x: x / args.t_model * (u.cm / u.day).to("km/s")) / 1e3
    z = cone["pos_x_min"].apply(lambda x: x / args.t_model * (u.cm / u.day).to("km/s")) / 1e3

    _surf = ax.scatter3D(x, y, z, c=-cone["fni"], cmap=plt.get_cmap("viridis"))

    # fig.colorbar(_surf, shrink=0.5, aspect=5)

    ax.set_xlabel(r"x [10$^3$ km/s]")
    ax.set_ylabel(r"y [10$^3$ km/s]")
    ax.set_zlabel(r"z [10$^3$ km/s]")

    # plt.scatter(cone[f'pos_x_min']/1e11, cone[f'pos_y_min']/1e11)
    plt.show()


def addargs(parser: argparse.ArgumentParser) -> None:
    parser.add_argument(
        "-modelpath",
        default=[],
        nargs="*",
        action=at.AppendPath,
        help="Path to ARTIS model folders with model.txt and abundances.txt",
    )

    parser.add_argument(
        "-sliceaxis",
        type=str,
        default="x",
        help=(
            "Choose x, y or z. The 1D model will be made based on this axis."
            "If cone then cone made around sliceaxis. Otherwise model along axis."
        ),
    )

    parser.add_argument(
        "--positive_axis", action="store", default=True, help="Make 1D model from positive axis. Default is True"
    )

    parser.add_argument(
        "--makefromcone",
        action="store",
        default=True,
        help="Make 1D model from cone around axis. Default is True.If False uses points along axis.",
    )


def main(args: argparse.Namespace | None = None, argsraw: t.Sequence[str] | None = None, **kwargs: t.Any) -> None:
    """Make 1D model from cone in 3D model."""
    if args is None:
        parser = argparse.ArgumentParser(
            formatter_class=at.CustomArgHelpFormatter,
            description=__doc__,
        )
        addargs(parser)
        parser.set_defaults(**kwargs)
        args = parser.parse_args(argsraw)

    if not args.modelpath:
        args.modelpath = [Path()]

    axes = ["x", "y", "z"]
    args.other_axis1 = next(ax for ax in axes if ax != args.sliceaxis)
    args.other_axis2 = next(ax for ax in axes if ax not in [args.sliceaxis, args.other_axis1])

    # remember: models before scaling down to artis input have x and z axis swapped compared to artis input files

    make_1D_model_files(args)

    # make_plot(args) # Uncomment to make 3D plot todo: add command line option


if __name__ == "__main__":
    main()
