"""Main functions of `estout` package"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.

# %% ../nbs/00_core.ipynb 4
from __future__ import annotations
from pathlib import Path 
from typing import List, Dict, Literal, Union, Callable
import importlib

import numpy as np
import pandas as pd
import statsmodels.api as sm
from linearmodels import PanelOLS

from . import utils 

# %% auto 0
__all__ = ['collect_stats', 'to_df', 'to_tex', 'to_pdf']

# %% ../nbs/00_core.ipynb 7
def collect_stats(res, # Results object to extract stats from
                  get_default_stats = True, # If True, returns all stats implemented by the `f'{package}_results'` module
                  add_stats: dict=None, # Keys are stats to extract in addition to the default ones; values are attributes of `res` or callables
                  add_literals: dict=None, # Additional info to be added to output dict; values must be scalars
                  ) -> dict:
    """Collects stats from `res` object. Stats in `add_stats` can override default stats"""

    out = {'package': res.__module__.split('.')[0]}
    results_module = importlib.import_module(f"estout.{out['package']}_results")

    if get_default_stats:
        for stat in results_module.__all__:
            out[stat] = utils.rgetattr(results_module, stat)(res)

    if add_stats is not None:
        for stat, val in add_stats.items():
            if isinstance(val, str): out[stat] = utils.rgetattr(res, val)
            elif isinstance(val, Callable): out[stat] = val(res)
            else: raise TypeError("Values of add_stats dict must be either strings or callables")

    if add_literals is not None:
        out.update(add_literals)
        
    return out

# %% ../nbs/00_core.ipynb 13
def to_df(res_list: List[dict], # List of outputs from `collect_stats()`
          which_xvars: list=None, # If None, report all xvars
          stats_body: list=['params', 'tstats'], # Each element of `res_list` needs to have these stats as keys; values must be pd.Series
          stats_bottom: list=['r2', 'nobs'], # Each element of `res_list` needs to have these stats as keys; values must be scalars
          labels: dict=None,
          add_formats: dict=None  
          ) -> pd.DataFrame: 
    """Combines results from multiple `collect_stats()` outputs into a single `pd`.DataFrame"""  
    
    formats = utils.default_formats()
    if add_formats is not None: formats.update(add_formats)
    
    columns = []
    for i,res in enumerate(res_list):
        newcol = pd.concat([res[x] for x in stats_body], axis=1, ignore_index=True).set_axis(stats_body, axis=1)
        for x in stats_body:
            newcol[x] = newcol[x].map(formats[x].format)
            if x == 'params':
                newcol[x] += utils.get_stars(res['pvalues'])
            else:
                newcol[x] = '(' + newcol[x] + ')'
        newcol = newcol.stack(level=0) #set_index('coeff_names')
        columns.append(newcol)

    out = pd.concat(columns, axis = 1)
    if which_xvars is not None: out = out.loc[which_xvars].copy()
    
    for i,res in enumerate(res_list):
        for x in stats_bottom:
            out.loc[x,i] = formats[x].format(res[x]) if x in formats else res[x]

    if labels is not None:
        for var in set(out.droplevel(1).index):
            if var in labels: out = out.rename(index={var:labels[var]}, level=0) 

    #Clean up row names 
    out = out.reset_index()
    out.loc[~out['level_1'].isin(['params','']), 'level_0'] = ''
    out = out.set_index(['level_0','level_1'])
    out.index.names = (None,None)

    return out.astype('string').fillna('')

# %% ../nbs/00_core.ipynb 16
def to_tex(dfs: pd.DataFrame|List[pd.DataFrame], # DataFrame(s) to be converted to tex table; if multiple, they will be panels in a larger table
            outfile: Path|str=None, # Where to save resulting tex output
            title: str='Table title', # Table title
            notes: str='Table description', # Some call this the table caption
            notes_on_top: bool=True, # Set to False if you want table description (caption) to be at the bottom
            label: str='', # Table label (for referencing within LaTex document)
            table_type: Literal['table','sidewaystable']='table',
            font_size: str='\\footnotesize', # Gets applied to the table contents as well as its caption
            addtocounter: int=0, # Set to -1 for tables that are just a continuation of a table on a new page
            
            panel_title: List[str]=None, # One element in the list for each dataframe in `dfs`
            palign: Literal['l','r','c']='l', # Alignment of panel title 
            col_groups: List[dict]=None, # Keys are group names; values are lists of consecutive indices of columns in the group
            col_names: List[Union[list,bool]]=True, # If False, none; if True, use df column names; if list, gives custom column names
            hlines: List[List[int]]=None, # Row indices under which to place hline
            tabular_env: str='tabular*' #LaTex tabular environment
            ) -> str:
    """Create tex code to generate table from one or more dataframes"""

    if isinstance(dfs, pd.DataFrame): dfs = [dfs]
    if panel_title is None: panel_title = ['']*len(dfs)
    if col_groups is None: col_groups = [None]*len(dfs)
    if col_names in [True, False]: col_names = [col_names]*len(dfs)
    if hlines is None: hlines = [[]]*len(dfs)
    body =  '\n \\smallskip \n'.join([utils.df_to_tex(dfs[i], panel_title=panel_title[i], 
                                                        palign=palign, tabular_env=tabular_env,
                                                        col_groups=col_groups[i], col_names=col_names[i], hlines=hlines[i]) 
                            for i in range(len(dfs))])
    
    pre = "\\newpage \n \\clearpage \n "
    pre += f"\\begin{{{table_type}}}[!h] {font_size} \n"
    pre += f"\\addtocounter{{table}}{{{addtocounter}}} \n"
    pre += f"\\caption{{\\textbf{{{title}}}}} \n"

    post = f"\\label{{{label}}} \n \\end{{{table_type}}} \n"

    notes_tex = ''
    if addtocounter==0: notes_tex = f"\\par {{{notes}}}"

    if notes_on_top:
        mid = notes_tex + f" \n\n \\vspace{{{'2mm'}}} \n\n {body} \n"
    else:
        mid = f"{body} \n\n \\vspace{{{'2mm'}}} \n\n " + notes_tex + ' \n'

    content = pre + mid + post

    if outfile:
        with open(outfile, "+w") as f:
            f.write(content) 
    return content

# %% ../nbs/00_core.ipynb 18
def to_pdf(outfile: str, # Path to .tex file where combined tables are saved (must contain .tex extension)
            table_tex_code: str|Path|List[Union[str,Path]]=None, # String(s) or Paths to files containing table tex code (e.g. like outputs of to_tex())
            article_spec=r'\documentclass[11pt]{article}',
            captionsetup="format=plain, labelsep=newline, labelfont = bf, justification=centering",
            make_pdf: bool=True, 
            open_pdf: bool=False):
    r"""Creates PDF with one or more tables given their tex code (from \being{table} to \end{table})"""

    if isinstance(table_tex_code, str): table_tex_code = [table_tex_code]
    tex_strings = []
    for t in table_tex_code:
        if isinstance(t, Path):
            with open(t,'r') as f:
                tex_strings.append(f.read())
        else: tex_strings.append(t)
    tables = ' \n '.join(tex_strings)

    preamble = r"\usepackage{booktabs,setspace,graphicx,epstopdf,tabularx, bigstrut,textcomp, outlines}" + "\n"
    preamble += r"\usepackage{amsmath,amsfonts,amssymb,amsthm,caption}" + "\n"
    preamble += r"\usepackage[margin=1in]{geometry}" + "\n"
    preamble += f"\\captionsetup{{{captionsetup}}}" + "\n"    

    content = "\n".join([article_spec, preamble, r'\begin{document}', tables, r'\end{document}'])

    with open(outfile, "w") as f:
        f.write(content) 
    if make_pdf:
        pdf_path = utils.make_pdf_from_tex(outfile)
        if open_pdf: return utils.open_pdf_file(pdf_path)
