import pandas as pd
from pathlib import Path


def _local_data_path(filename='morph_data', ext="csv"):
  """
  Build a path like: <this_file_dir>/data/<filename>.<ext>

  Parameters
  ----------
  filename : str
      Base filename (without extension)
  ext : str, default "csv"
      File extension (without the dot)

  Returns
  -------
  pathlib.Path
      Full path to the data file
  """
  work_dir = Path(__file__).resolve().parent
  return work_dir / f"{filename}.{ext}"


def get_data(key):
  """
  Retrieve a dataset entry using a key-path of the form
  "<brain region>/<neuron class>/<subcellular section>".

  The argument `data_path` is interpreted as a slash-separated path of keys used
  to traverse a nested dataset dictionary. The selected dataset is expected to
  contain both Sholl-plot statistics and bifurcation statistics; when both are
  available, this function returns a standardized dictionary compatible with
  `compute_rates`.

  Parameters
  ----------
  key : str
      Dataset identifier expressed as a key path:

      "<brain region>/<neuron class>/<subcellular section>"

      Examples:
      - "CTX/pyr/apical"
      - "HPC/pyr/basal"

      Each component is used as a successive key lookup into the nested dataset
      container.

  Returns
  -------
  dict
      If both Sholl and bifurcation information are present for the selected dataset,
      returns:

      data = {
        "sholl": {
          "bin_size": float,
          "mean": numpy.ndarray,   # shape (K,)
          "var":  numpy.ndarray,   # shape (K,)
        },
        "bifurcations": {
          "mean": float,
          "var":  float,
        },
      }

      Where:
      - `data["sholl"]["bin_size"]` is the spatial bin size used to define Sholl shells
      - `data["sholl"]["mean"]` is the mean Sholl intersection count per radial bin
      - `data["sholl"]["var"]` is the variance of the Sholl intersection count per bin
      - `data["bifurcations"]["mean"]` is the mean bifurcation count
      - `data["bifurcations"]["var"]` is the variance of the bifurcation count

  Raises
  ------
  KeyError
      If any key along `data_path` is missing (brain region, neuron class, or section)
  ValueError
      If the selected dataset does not contain both Sholl and bifurcation data, or
      if the provided arrays have incompatible shapes

  Notes
  -----
  - `data_path` is a *key path*, not a filesystem path
  - The function assumes the dataset entry referenced by `data_path` includes:
      - Sholl bin size, mean array, variance array
      - Bifurcation mean and variance

  Examples
  --------
  >>> data = get("CTX/pyr/apical")
  >>> data["sholl"]["bin_size"]
  50.0
  >>> data["bifurcations"]["mean"]
  12.3
  """
  data = {}

  # split the key
  parts = tuple(p.strip() for p in key.split("/") if p.strip())
  if len(parts) != 2:
      raise ValueError(f"Expected key like 'area/neuron_type', got: {key!r}")
  area, neuron_type = parts
    
  # load data
  df = pd.read_csv(_local_data_path(), index_col=0)

  # select specific area and neuron type
  df = df[(df['area'] == area) & (df['neuron_type'] == neuron_type)]
  
  # neuron name unnecessary
  df.drop(['area', 'neuron_type', 'neuron_name'], axis=1, inplace=True)

  # statistics
  df = df.groupby('section_type').describe()

  # select only a subset of columns
  df = df.loc[:, df.columns.get_level_values(1).isin(['mean', 'std', 'min', 'max'])]

  # get subsections
  for section_type, row in df.iterrows():
    data[section_type] = {}

    print()

    # get statistics
    for data_type in ['bifurcation_count', 'total_length']:
      tmp = row.loc[row.index.get_level_values(0) == data_type, :]
      tmp.index = tmp.index.droplevel(0)
      data[section_type][data_type] = tmp.to_dict()

    # count neurites at the soma
    tmp = row.loc[row.index.get_level_values(0) == 'Count0', :]
    tmp.index = tmp.index.droplevel(0)
    data[section_type]['primary_count'] = tmp.to_dict()
    
    # sholl plots
    tmp = row.loc[row.index.get_level_values(0).str.startswith('Count'), :]
    data[section_type]['sholl_plot'] = {
      'bin_size':row[('bin_size', 'mean')].tolist(),
      'mean':tmp.loc[tmp.index.get_level_values(1) == 'mean', :].tolist(),
      'std':tmp.loc[tmp.index.get_level_values(1) == 'std', :].tolist()
      }

  return data

