import json
from itertools import chain
from collections.abc import Iterable, Mapping
from typing import Optional, List, Dict, Union, Literal

import numpy as np
import pandas as pd


def segment_properties_json(
    df: pd.DataFrame,
    label_col: Optional[str] = None,
    *,
    description_col: Optional[str] = None,
    string_cols: List[str] = [],
    number_cols: List[str] = [],
    tag_cols: List[str] = [],
    tag_prefix_mode: Literal['all', 'disambiguate', None] = 'all',
    sort_tags: bool = False,
    tag_descriptions: Dict[str, str] = {},
    col_descriptions: Dict[str, str] = {},
    drop_empty: bool = True,
    output_path: Optional[str] = None,
) -> Dict[str, Union[str, Dict]]:
    """
    Given a DataFrame, construct the JSON representation for neuroglancer segment properties
    (i.e. segment properties 'info' file) according to the neuroglancer spec:
    https://github.com/google/neuroglancer/blob/master/src/datasource/precomputed/segment_properties.md

    Args:
        df:
            DataFrame whose index contains segment IDs.
            For clarity of intent, we require that the Index be named 'segment' or 'body'.

            All columns will be converted to segment properties. If a column is not
            explicitly listed in the following arguments, its property type will be inferred
            from the column dtype (or column name, in the case of label/description).

            The column dtypes can be string, category, or number, but 64-bit values
            will be downcast to 32-bit. Boolean columns are only valid as tags.

            The order of the input columns determines the order in which properties will be written.
            Neuroglancer preserves that order when displaying columns of numeric properties.

        label_col:
            Name of the column to use for the 'label' property, which is shown in neuroglancer by default.
            If the column happens to be named 'label', then you need not provide it here.
            Or, if your dataframe contains only one column (and it's non-numeric), we assume it is the label.

        description_col:
            Name of the column to use as the 'description' property.
            If you name the column 'description', then you need not provide it here.

        string_cols:
            Names of the columns to represent as 'string' properties.
            (Usually unnecessary.  Non-numeric columns become string properties by default,
            unless they are named in tag_cols.)

        number_cols:
            Names of the columns to represent as 'number' properties.
            (Usually unnecessary.  Numeric columns become number properties by default.)

        tag_cols:
            Names of the columns from which to generate the (combined) 'tags' property.
            If you want a column to be used for both tags _and_ a different property,
            be sure to list it explicitly in both arguments.  For example:

            .. code-block:: python

                segment_properties_json(df, label_col='cell_class', tag_cols=['cell_class', ...])

            In addition to string/category inputs, columns with dtype bool can also be
            used for tags, in which case the False items are discarded and the True items
            are tagged with the column name.

        tag_prefix_mode:
            Either ``'all'`` or ``'disambiguate'`` or ``None``.

            - If ``'all'``, then all tags will be prefixed with the name of their source column,
              (e.g. ``status:Anchor``), other than those which came from boolean columns.

            - If ``'disambiguate'``, then only tags which appear in multiple columns
              be prefixed with the name of their source column.

            - If None, then no disambiguation is performed.

        sort_tags:
            If True, order tags alphabetically (after prefixes are applied, if any) so they
            are listed in sorted order in the neuroglancer UI.
            Otherwise, tags are listed according to the order given by 'tag_cols',
            with tags from earlier columns preceding tags from later columns,
            but tags within each column will be sorted with respect to each other.
            Note: For categorical input columns with ordered categories, the sort order
            is determined by the categorical's intrinsic ordering, not alphanumeric sorting.

        tag_descriptions:
            Optional. A dict of {tag: description} describing each tag value.
            Missing tags will be given a default description.

        col_descriptions:
            Optional. A dict of {column: description} describing each property
            (input column) other than 'tags'.

        drop_empty:
            If any IDs in the input have no non-empty (null or "") properties,
            then drop them from the output entirely so they don't show up in
            neuroglancer's default segment list.

        output_path:
            If provided, export the JSON to a file.

    Returns:
        JSON data (as a dict) which should be written into an 'info' file to
        host neuroglancer precomputed segment properties.

    Example:

        .. code-block:: ipython

            In [41]: print(df)
            Out[41]:
                        type              class       hemilineage  has_soma  PreSyn  PostSyn
            segment
            910719  IN13B018   intrinsic_neuron               13B      True     786     3164
            160503    PFNm/p            central          DM3_CX_p      True      40      177
            908822  IN03B079   intrinsic_neuron               03B      True     176      951
            10552      DNp29  descending_neuron  putative_primary      True    1567     4559
            43797     MeVP11  visual_projection               NaN      True     158      504
            232845     R1-R6            sensory               NaN     False      25       18
            547654     R1-R6            sensory               NaN     False      57       29
            803197  IN07B019   intrinsic_neuron               TBD      True     796     3615
            75597      KCg-m            central              MBp4      True     223      802
            165485   vDeltaK            central               NaN      True      36      282

            In [42]: info = segment_properties_json(df, 'type', tag_cols=['class', 'hemilineage', 'has_soma'])

            In [43]: # Note: The output below has been indented by hand for clarity.
                ...: print(json.dumps(info))
            {
                "@type": "neuroglancer_segment_properties",
                "inline": {
                    "ids": ["910719", "160503", "908822", "10552", "43797", "232845", "547654", "803197", "75597", "165485"],
                    "properties": [
                        {
                            "id": "type",
                            "type": "label",
                            "values": ["IN13B018", "PFNm/p", "IN03B079", "DNp29", "MeVP11", "R1-R6", "R1-R6", "IN07B019", "KCg-m", "vDeltaK"]
                        },
                        {
                            "id": "PreSyn",
                            "type": "number",
                            "data_type": "int32",
                            "values": [786, 40, 176, 1567, 158, 25, 57, 796, 223, 36]
                        },
                        {
                            "id": "PostSyn",
                            "type": "number",
                            "data_type": "int32",
                            "values": [3164, 177, 951, 4559, 504, 18, 29, 3615, 802, 282]
                        },
                        {
                            "id": "tags",
                            "type": "tags",
                            "tags": [
                                "class:central", "class:descending_neuron", "class:intrinsic_neuron",
                                "class:sensory", "class:visual_projection",
                                "has_soma",
                                "hemilineage:03B", "hemilineage:13B", "hemilineage:DM3_CX_p", "hemilineage:MBp4",
                                "hemilineage:TBD", "hemilineage:putative_primary"
                            ],
                            "values": [
                                [2, 5, 7],
                                [0, 5, 8],
                                [2, 5, 6],
                                [1, 5, 11],
                                [4, 5],
                                [3],
                                [3],
                                [2, 5, 10],
                                [0, 5, 9],
                                [0, 5]
                            ]
                        }
                    ]
                }
            }
    """
    df, string_cols, number_cols, tag_cols = _validate_args(
        df, label_col, description_col, string_cols, number_cols,
        tag_cols, tag_prefix_mode, tag_descriptions, col_descriptions
    )

    if drop_empty:
        df = _drop_empty_rows(df)

    # Determine the neuroglancer property type for each column (except tags).
    scalar_types = _scalar_property_types(
        df, label_col, description_col, string_cols, number_cols, tag_cols
    )

    # Construct JSON for all properties except tags.
    json_props = []
    for col, prop_type in scalar_types.items():
        jsn = _scalar_property_json(df[col], prop_type, col_descriptions.get(col))
        json_props.append(jsn)

    # Construct JSON for tags.
    if len(tag_cols):
        jsn = _tags_property_json(df, tag_cols, tag_prefix_mode, sort_tags, tag_descriptions)
        json_props.append(jsn)

    # Assemble the complete output.
    segment_ids = [str(idx) for idx in df.index]
    info = {
        '@type': 'neuroglancer_segment_properties',
        'inline': {
            'ids': segment_ids,
            'properties': json_props
        }
    }

    if output_path:
        with open(output_path, 'w') as f:
            json.dump(info, f)

    return info


def _validate_args(df, label_col, description_col, string_cols, number_cols,
                   tag_cols, tag_prefix_mode, tag_descriptions, col_descriptions):
    """
    Basic checks and convenience conversions where appropriate.
    (Series -> DataFrame, str -> list)
    """
    if isinstance(df, pd.Series):
        df = df.to_frame()

    if isinstance(string_cols, str):
        string_cols = [string_cols]
    if isinstance(number_cols, str):
        number_cols = [number_cols]
    if isinstance(tag_cols, str):
        tag_cols = [tag_cols]

    assert isinstance(df, pd.DataFrame)
    assert isinstance(string_cols, Iterable)
    assert isinstance(number_cols, Iterable)
    assert isinstance(tag_cols, Iterable)
    assert isinstance(tag_descriptions, Mapping)
    assert isinstance(col_descriptions, Mapping)

    assert df.index.name in ('body', 'segment'), \
        f"Index name must be 'body' or 'segment', not '{df.index.name}'"
    assert tag_prefix_mode in ('all', 'disambiguate', None), \
        f"Invalid tag_prefix_mode: {tag_prefix_mode}."
    assert not (dupes := df.columns.duplicated()).any(), \
        f"Duplicated column names: {df.columns[dupes].tolist()}"

    # Check for columns that were listed in multiple arguments (except tag_cols).
    listed_scalar_cols = [label_col, description_col, *string_cols, *number_cols]
    listed_scalar_cols = [*filter(None, listed_scalar_cols)]
    listed_scalar_cols = pd.Series(listed_scalar_cols)
    dupes = listed_scalar_cols.loc[listed_scalar_cols.duplicated()].unique()
    if dupes := sorted(dupes):
        raise RuntimeError(f"Some columns were included in multiple arguments: {dupes}")

    return df, string_cols, number_cols, tag_cols


def _drop_empty_rows(df):
    """
    Drop rows which contain only NaN, "", and False, and thus
    will have no non-empty value for any neuroglancer property.
    """
    if df.empty:
        return df

    # Note that the only columns which are permitted to
    # have dtype 'bool' are the boolean tag cols (if any).
    # Therefore, False is considered 'empty'.
    bool_cols = df.select_dtypes(include=bool).columns
    other_cols = df.columns.difference(bool_cols)

    valid_bool = valid_other = False
    if len(bool_cols) > 0:
        valid_bool = df[bool_cols].any(axis=1)
    if len(other_cols) > 0:
        not_null = df[other_cols].notnull()
        not_empty = df[other_cols] != ''
        valid_other = (not_null & not_empty).any(axis=1)

    return df.loc[valid_bool | valid_other]


def _scalar_property_types(df, label_col, description_col, string_cols, number_cols, tag_cols):
    """
    Determine the names and neuroglancer property types of all scalar (non-tag)
    properties that can be extracted from the given DataFrame.

    The property types will be determined from the user's explicitly provided lists
    when available; otherwise they're inferred from the column name/dtype.

    Here, 'scalar' includes all non-tag types: number, string, label, description.
    """
    # Start with the user-provided types and then auto-type the other columns below.
    # Although we won't include tag columns in the result, we temporarily store
    # them in prop_types to avoid auto-typing them.
    # Note:
    #   The user is permitted to list tag columns in the scalar columns too,
    #   in which case the scalar type takes precedence in this function.
    prop_types = {c: 'tags' for c in tag_cols}
    prop_types |= {c: 'string' for c in string_cols}
    prop_types |= {c: 'number' for c in number_cols}

    if label_col:
        prop_types[label_col] = 'label'
    if description_col:
        prop_types[description_col] = 'description'

    if unknown_cols := set(prop_types.keys()) - set(df.columns):
        raise RuntimeError(f"Columns not found: {unknown_cols}")

    # For convenience, if there's only one column, and it wasn't explicitly listed,
    # AND it's non-numeric, we presume it's meant to be the 'label' column.
    if (
        len(df.columns) == 1
        and not prop_types
        and (df.dtypes.iloc[0] == 'string' or not np.issubdtype(df.dtypes.iloc[0], np.number))
    ):
        prop_types = {df.columns[0]: 'label'}

    # Infer the types of unlisted columns from either the name or dtype
    for name, dtype in df.dtypes.items():
        if dtype in (bool, 'boolean') and prop_types.get(name) != 'tags':  # noqa: E721
            raise RuntimeError(f"Column '{name}': Boolean columns are only valid as tag_cols")
        elif prop_types.get(name) == 'number' and not np.issubdtype(dtype, np.number):
            raise RuntimeError(f"Column '{name}': Not valid as number_cols (dtype: {dtype})")
        elif name in prop_types:
            continue
        elif name == 'label':
            prop_types['label'] = 'label'
        elif name == 'description':
            prop_types['description'] = 'description'
        elif dtype in ("object", "string"):
            prop_types[name] = 'string'
        elif np.issubdtype(dtype, np.number):
            prop_types[name] = 'number'
        else:
            prop_types[name] = 'string'

    # Re-order to match original input columns.
    # Property order determines appearance order in neuroglancer,
    # at least for numeric properties.
    prop_types = {c: prop_types[c] for c in df.columns}

    # Return scalar properties only
    prop_types = {k:v for k,v in prop_types.items() if v != 'tags'}
    return prop_types


def _scalar_property_json(s, prop_type, description):
    """
    Constructs the JSON for any segment property
    other than the 'tags' property.
    """
    if prop_type == 'number':
        return _scalar_number_property_json(s, description)

    prop = {
        'id': s.name,
        'type': prop_type,
        'values': s.fillna("").astype(str).tolist()
    }

    if description:
        prop['description'] = description

    return prop


def _scalar_number_property_json(s, description):
    """
    Constructs the JSON for a 'number' property.
    """
    dtype_name = s.dtype.name

    if s.dtype in (np.int64, np.uint64):
        dtype_name = _select_int64_downcast(s)

    if np.issubdtype(s.dtype, np.floating):
        dtype_name = 'float32'
        if s.isnull().any():
            raise RuntimeError(f"Numeric column '{s.name}' contains NaN.")

    prop = {
        'id': s.name,
        'type': 'number',
        'data_type': dtype_name,
        'values': s.tolist()
    }

    if description:
        prop['description'] = description

    return prop


def _select_int64_downcast(s):
    """
    Return either "int32" or "uint32" if the given (u)int64 Series
    can be safely converted to one of those dtypes without overflow.
    Otherwise, raise an error.
    """
    for dtype32 in (np.int32, np.uint32):
        info32 = np.iinfo(dtype32)
        if s.min() >= info32.min and s.max() <= info32.max:
            return info32.dtype.name

    raise RuntimeError(
        f"Can't create a property for column: '{s.name}'. "
        "Neuroglancer doesn't support 64-bit integer properties, "
        "and your data exceeds the limits of (u)int32."
    )


def _tags_property_json(df, tag_cols, tag_prefix_mode, sort_tags, tag_descriptions):
    """
    Constructs the JSON for the 'tags' segment property.
    """
    tags_df = df[[]].copy()

    # Individually convert each column to Categorical
    # before we combine into a unified Categorical below.
    for c in tag_cols:
        tags_df[c] = _convert_to_categorical(df[c])

    _insert_tag_prefixes(tags_df, tag_prefix_mode, df.dtypes)

    # Unified tag list: deduplicate, but preserve order by default.
    tag_sets = (dtype.categories for dtype in tags_df.dtypes)
    unique_tags = pd.Series([*chain(*tag_sets)]).unique().tolist()
    if sort_tags:
        unique_tags = sorted(unique_tags)

    # Convert all columns to a single unified Categorical dtype.
    tags_df = tags_df.astype(pd.CategoricalDtype(categories=unique_tags))

    # Tags are represented in JSON as a list-of-lists of sorted codes.
    codes_df = pd.DataFrame({c: s.cat.codes for c, s in tags_df.items()})
    sorted_codes = np.sort(codes_df.to_numpy(), axis=1).tolist()

    # Drop nulls
    for row in sorted_codes:
        row[:] = [x for x in row if x != -1]

    prop = {
        'id': 'tags',
        'type': 'tags',
        'tags': unique_tags,
        'values': sorted_codes,
    }

    if tag_descriptions:
        prop['tag_descriptions'] = _tag_description_list(unique_tags, tag_descriptions)

    return prop


def _convert_to_categorical(s):
    """
    Convert the given Series to a Categorical Series suitable for tags.
    """
    if s.dtype == 'category':
        s = s.cat.remove_unused_categories()

    if s.dtype in (bool, 'boolean'):
        s = s.astype('category', copy=False)
        s = s.cat.rename_categories({True: s.name})
        if False in s.dtype.categories:
            s = s.cat.remove_categories([False])

    s = s.astype('category', copy=False)
    s = _replace_spaces(s)

    # We interpret empty string as null
    if '' in s.dtype.categories:
        s = s.cat.remove_categories([''])

    return s


def _replace_spaces(s):
    """
    Replace all spaces in the given Categorical Series with underscores.
    (Spaces are forbidden in neuroglancer tags.)
    """
    renames = {
        cat: str(cat).replace(' ', '_')
        for cat in s.dtype.categories
    }

    try:
        return s.cat.rename_categories(renames)
    except ValueError:
        # If both "foo bar_baz" and "foo_bar baz" exist in the
        # original data, they will both map to "foo_bar_baz", and
        # rename_categories() complains about the duplicate category.
        return s.str.replace(' ', '_').astype('category')


def _insert_tag_prefixes(df, tag_prefix_mode, orig_dtypes):
    """
    Insert prefixes onto tags in df according to the tag_prefix_mode (if any).
    The columns of df must already have Categorical dtypes.

    Tags in columns which were originally bool (before we converted them
    to Categorical) require no prefix, so we refer to orig_dtypes to skip
    those columns.

    Modifies df in-place.
    """
    if tag_prefix_mode is None:
        return

    if tag_prefix_mode == 'disambiguate':
        _disambiguate_tags(df)
        return

    assert tag_prefix_mode == 'all'
    for c, s in list(df.items()):
        if orig_dtypes[c] in ('bool', 'boolean'):
            # Boolean columns need no prefix;
            # the column name is the entire tag.
            continue

        # Prefix must not contain spaces (forbidden by neuroglancer)
        # or colons (because we use ':' as the separator).
        prefix = s.name.replace(' ', '_').replace(':', '_')
        df[c] = s.cat.rename_categories([
            f'{prefix}:{cat}'
            for cat in s.dtype.categories
        ])


def _disambiguate_tags(df):
    """
    Given a dataframe in which all columns are Categoricals,
    find category values that are common across multiple columns
    and prepend a prefix (the column name) to such values
    to make sure no category value is duplicated from one column
    to the next.

    Modifies df in-place.
    """
    col_tags = [
        (col, tag)
        for col in df.columns
        for tag in df[col].dtype.categories
    ]
    col_tags = pd.DataFrame(col_tags, columns=['col', 'tag'])

    # Drop everything except the duplicate tags.
    is_dup = col_tags.duplicated('tag', keep=False)
    col_tags = col_tags.loc[is_dup].copy()

    # Prefix must not contain spaces (forbidden by neuroglancer)
    # or colons (because we use ':' as the separator).
    prefixes = col_tags['col'].str.replace(r'[: ]', '_')
    col_tags['new_tag'] = prefixes + ':' + col_tags['tag']

    for col, old_new_df in col_tags.groupby('col')[['tag', 'new_tag']]:
        renames = dict(old_new_df.to_numpy())
        df[col] = df[col].cat.rename_categories(renames)


def _tag_description_list(unique_tags, tag_descriptions):
    """
    Given the list of all tags and a mapping of tags to descriptions,
    return the descriptions in the same order as unique_tags.
    """
    tag_descriptions = {
        str(k).replace(' ', '_'): v
        for k,v in tag_descriptions.items()
    }

    td = []
    for t in unique_tags:
        d = tag_descriptions.get(t, None)

        # If we didn't find it, try stripping the tag prefix (if any)
        if not d and ':' in t:
            t2 = t[1+t.index(':'):]
            d = tag_descriptions.get(t2, None)

        # If we still didn't find it, emit the tag itself.
        td.append(d or t)

    return td


def segment_properties_to_dataframe(js: dict, consolidate_tags_by_prefix: bool = True, return_separate_tags: bool = False):
    """
    Converts segment properties info JSON to DataFrame.

    Args:
        js:
            The JSON representation of the segment properties.

        consolidate_tags_by_prefix:
            If True, consolidate boolean tag columns into Categorical columns
            according to their prefix, assuming the data was written
            using according to the naming conventions followed
            by segment_properties_json(..., tag_prefix_mode="all").
            Otherwise, a separate boolean column is created for each
            unique tag -- potentially tons of columns (and RAM)!

        return_separate_tags:
            If True, return two dataframes, for scalar columns and tag columns.
            Otherwise, combine them into a single dataframe.

    Returns:
        If return_separate_tags=True, return (scalar_df, tags_df).
        By default, concatenate them across columns into a single dataframe.
    """
    segment_ids = [*map(int, js['inline']['ids'])]
    segment_ids = pd.Index(segment_ids, name='segment')

    all_props = js['inline']['properties']
    scalar_props = [prop for prop in all_props if prop['type'] != 'tags']
    tags_props =   [prop for prop in all_props if prop['type'] == 'tags']

    scalar_values = {
        prop['id']: prop['values']
        for prop in scalar_props
    }

    scalar_dtypes = {
        prop['id']: prop['data_type']
        for prop in scalar_props
        if 'data_type' in prop
    }

    scalar_df = pd.DataFrame(scalar_values, segment_ids).astype(scalar_dtypes)

    if tags_props:
        tags_df = _convert_tags_to_dataframe(segment_ids, tags_props, consolidate_tags_by_prefix)
    else:
        tags_df = scalar_df[[]]

    if return_separate_tags:
        return scalar_df, tags_df

    # It's possible for a scalar column to also be used in the tags.
    # We'll drop the scalar and just keep the one from the tags,
    # since the tag column takes less RAM.
    scalar_df = scalar_df[[c for c in scalar_df.columns if c not in tags_df.columns]]
    return pd.concat((scalar_df, tags_df), axis=1)


def _convert_tags_to_dataframe(segment_ids, tags_props, consolidate_tags_by_prefix):
    code_lists = tags_props[0]['values']
    unique_tags = tags_props[0]['tags']

    segment_count = len(code_lists)
    tag_count = len(unique_tags)

    # Flatten the lists of codes,
    # but keep track of the rows they came from.
    flat_tag_codes = pd.Series([*chain(*code_lists)])
    code_rows = np.repeat(
        range(segment_count),
        [*map(len, code_lists)]
    )

    if not consolidate_tags_by_prefix:
        # Return a boolean column for every unique tag.
        # Could be tons of columns!
        flags_matrix = np.zeros((segment_count, tag_count), bool)
        flags_matrix[code_rows, flat_tag_codes] = True
        return pd.DataFrame(flags_matrix, segment_ids, unique_tags)

    tag_prefixes = pd.Series([tag.split(':', 1)[0] for tag in unique_tags]).astype('category')
    flat_prefixes = flat_tag_codes.map(tag_prefixes)

    # Assemble the flattened codes into a matrix (row per segment, column per prefix).
    # Default is -1 (categorical code for a missing value).
    code_matrix = -1 * np.ones((segment_count, pd.Series(tag_prefixes).nunique()), np.int32)
    code_matrix[code_rows, flat_prefixes.cat.codes] = flat_tag_codes

    tag_categoricals = {}
    for prefix, code_column in zip(tag_prefixes.cat.categories, code_matrix.T):
        # Initialize as a categorical that corresponds to all tags combined,
        # then reduce to a categorical for just this column's values (without prefixes).
        tag_cat = pd.Categorical.from_codes(code_column, unique_tags)
        tag_cat = tag_cat.remove_unused_categories()
        tag_cat = tag_cat.rename_categories({
            tag: tag.split(':', 1)[1] if ':' in tag else tag
            for tag in tag_cat.dtype.categories
        })
        tag_categoricals[prefix] = tag_cat

    return pd.DataFrame(tag_categoricals, segment_ids)
