import re
from io import StringIO
from pathlib import Path
from typing import Sequence

import astropy
import numpy as np
import pandas as pd
from astropy.io import fits
from astropy.table import Table


def load_tdat(path: str | Path, columns: Sequence[str] | None = None) -> pd.DataFrame:
  """
  Loads a tdat table. Transportable Database Aggregate Table (TDAT) is a data
  structure created by NASA's Heasarc project. For more information, see [TDAT]_

  Parameters
  ----------
  path : str | Path
    Path to the table to read
  columns : Sequence[str] | None
    If specified, only the column names in list will be loaded. Can be used to
    reduce memory usage

  Returns
  -------
  pd.DataFrame
    The table as a pandas dataframe
    
  References
  ----------
  .. [TDAT] Transportable Database Aggregate Table (TDAT) Format.
      `<https://heasarc.gsfc.nasa.gov/docs/software/dbdocs/tdat.html>`_
  """
  path = Path(path)
  content = path.read_text()
  header = re.findall(r'line\[1\] = (.*)', content)[0].replace(' ', '|')
  data = content.split('<DATA>\n')[-1].split('<END>')[0].replace('|\n', '\n')
  tb = header + '\n' + data
  df = pd.read_csv(StringIO(tb), sep='|')
  
  if columns:
    df = df[columns]
  return df


def load_fits(path: str | Path, columns: Sequence[str] | None = None) -> pd.DataFrame:
  """
  Loads a fits table

  Parameters
  ----------
  path : str | Path
    Path to the table to read
  columns : Sequence[str] | None
    If specified, only the column names in list will be loaded. Can be used to
    reduce memory usage

  Returns
  -------
  pd.DataFrame
    The table as a pandas dataframe
  """
  with fits.open(path) as hdul:
    table_data = hdul[1].data
    table = Table(data=table_data, names=columns)
    return table.to_pandas()


def load_csv(path: str | Path, columns: Sequence[str] | None = None) -> pd.DataFrame:
  """
  Loads a csv table. ASCII tables with columns delimited by a comma.

  Parameters
  ----------
  path : str | Path
    Path to the table to read
  columns : Sequence[str] | None
    If specified, only the column names in list will be loaded. Can be used to
    reduce memory usage

  Returns
  -------
  pd.DataFrame
    The table as a pandas dataframe
  """
  return pd.read_csv(path, names=columns)


def load_tsv(path: str | Path, columns: Sequence[str] | None = None) -> pd.DataFrame:
  """
  Loads a tsv table. ASCII tables with columns delimited by a white character,
  e.g. space or tab.

  Parameters
  ----------
  path : str | Path
    Path to the table to read
  columns : Sequence[str] | None
    If specified, only the column names in list will be loaded. Can be used to
    reduce memory usage

  Returns
  -------
  pd.DataFrame
    The table as a pandas dataframe
  """
  return pd.read_csv(path, delim_whitespace=True, names=columns)


def load_parquet(path: str | Path, columns: Sequence[str] | None = None) -> pd.DataFrame:
  """
  Loads a parquet table

  Parameters
  ----------
  path : str | Path
    Path to the table to read
  columns : Sequence[str] | None
    If specified, only the column names in list will be loaded. Can be used to
    reduce memory usage

  Returns
  -------
  pd.DataFrame
    The table as a pandas dataframe
  """
  return pd.read_parquet(path, columns=columns)


def load_feather(path: str | Path, columns: Sequence[str] | None = None) -> pd.DataFrame:
  """
  Loads a feather table

  Parameters
  ----------
  path : str | Path
    Path to the table to read
  columns : Sequence[str] | None
    If specified, only the column names in list will be loaded. Can be used to
    reduce memory usage

  Returns
  -------
  pd.DataFrame
    The table as a pandas dataframe
  """
  return pd.read_feather(path, columns=columns)


def load_table(path: str | Path, columns: Sequence[str] | None = None) -> pd.DataFrame:
  """
  This function tries to detect the table type comparing the file extension and
  returns a pandas dataframe of the loaded table.
  
  Supported table types:
  
  - Fits tables: .fit, .fits, .fz
  - ASCII tables: .csv, .tsv, .dat
  - Heasarc tables: .tdat
  - Arrow tables: .parquet, .feather

  Parameters
  ----------
  path : str | Path
    Path to the table to read
  columns : Sequence[str] | None
    If specified, only the column names in list will be loaded. Can be used to
    reduce memory usage

  Returns
  -------
  pd.DataFrame
    The table as a pandas dataframe

  Raises
  ------
  ValueError
    Raises an error if the file extension can not be detected
  """
  func_map = {
    '.fit': load_fits,
    '.fits': load_fits,
    '.fz': load_fits,
    '.csv': load_csv,
    '.tsv': load_tsv,
    '.dat': load_tsv,
    '.parquet': load_parquet,
    '.feather': load_feather,
    '.tdat': load_tdat,
  }
  
  path = Path(path)
  load_func = func_map.get(path.suffix)
  if load_func is None:
    raise ValueError(
      'Can not infer the load function for this table based on suffix. '
      'Please, use a specific loader.'
    )
  return load_func(path, columns=columns)


def save_table(data: pd.DataFrame, path: str | Path):
  path = Path(path)
  if path.suffix in ('.fit', '.fits'):
    Table.from_pandas(data).write(path, overwrite=True)
  elif path.suffix == '.csv':
    data.to_csv(path, index=False)
  elif path.suffix == '.parquet':
    data.to_parquet(path, index=False)
  elif path.suffix == '.dat':
    data.to_csv(path, sep=' ')
  elif path.suffix == '.tsv':
    data.to_csv(path, sep='\t')