"""Base classes, functions and other objects used across the package."""

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

# %% auto 0
__all__ = ['CODE_ROOT', 'PACKAGE_ROOT', 'ProjectFileSystem', 'SqliteDatabase', 'JsonDict', 'JsonFileReader',
           'list_available_devices', 'TextFileBaseReader', 'TextFileBaseIterator']

# %% ../nbs-dev/00_core.ipynb 4
import json
import os
import pandas as pd
import re
import sqlite3
import sys
import torch
import warnings
from configparser import ConfigParser
from eccore.core import validate_path, safe_path
from IPython.display import display, Markdown, HTML
from pathlib import Path
from pprint import pprint
from sqlite3 import Connection, Cursor
from typing import Any, Optional, Literal

try: from google.colab import drive
except: pass

# %% ../nbs-dev/00_core.ipynb 5
# Retrieve the package root
from . import __file__
CODE_ROOT = Path(__file__).parents[0]
PACKAGE_ROOT = Path(__file__).parents[1]

# %% ../nbs-dev/00_core.ipynb 10
class ProjectFileSystem:
    """Represent a project file system, return paths to key directories, provide methods to manage the file system.

    - Paths to key directories are based on whether the code is running locally or in the cloud.
    - First time it is used on a local computer, it must be registered as local and a project root path must be set.
    - A user configuration file is created in the user's home directory to store the project root path and whether the machine is local or not.

    > Technical note: `ProjectFileSystem` is a simpleton class
    """

    _instance = None
    _config_dir = '.metagentorch'
    _config_fname = 'metagentorch.cfg'
    _shared_project_dir = 'Metagenomics'
    
    def __new__(cls, *args, **kwargs):
        # Create instance if it does not exist yet
        if cls._instance is None:
            cls.home = Path.home().resolve()
            if kwargs.get('config_fname', None) is not None:
                cls._p2config = kwargs['config_fname']
            else:
                cls._p2config = cls.home / cls._config_dir / cls._config_fname
            cls._instance = super().__new__(cls)
        return cls._instance
    
    def __init__(
        self, 
        mount_gdrive:bool=True,       # True to mount Google Drive if running on Colab
        project_file:Path|None=None,  # Path to the project file. If None, use the one saved in the config file
        config_fname:Path|None=None,  # Path to a configuration file. If None, use the default one in the user's home directory
        ) -> None:

        # Discover where the script is being run
        self.is_colab = 'google.colab' in sys.modules       
        if self.is_colab and mount_gdrive:
            drive.mount('/content/gdrive')
            self.gdrive = Path('/content/gdrive/MyDrive')
        self.is_kaggle = 'kaggle_web_client' in sys.modules
        if self.is_kaggle:
            raise NotImplementedError(f"ProjectFileSystem is not implemented for Kaggle yet")
        if not self.is_colab and not self.is_kaggle and not self.is_local:
            msg = """
                    Code does not seem to run on the cloud but computer is not registered as local
                    If you are running on a local computer, you must register it as local by running
                    `ProjectFileSystem().register_as_local()`
                    before you can use the ProjectFileSystem class.
                    """
            warnings.warn(msg, UserWarning)

        # Set key directory paths
        self._project_root = Path()
        data_dir = 'data'
        if self.is_local or os.getenv("GITHUB_ACTIONS") == "true":
            cfg = self.read_config()
            path_str = cfg.get('Infra', 'project_root', fallback=None)
            data_dir = cfg.get('Infra', 'data_dir', fallback='data')
            if path_str is None: 
                msg = """
                Project root is not yet set in config file.
                To set it, use `ProjectFileSystem().set_project_root()`.
                """
                warnings.warn(msg)
            else:
                self._project_root = Path(path_str)
        self._data = self.project_root / data_dir
        self._nbs = self.project_root / 'nbs'

    def __call__(self) -> bool: return self.is_local

    def info(self) -> None:
        """Print basic info on the file system and the device"""
        print(f"Running {self.os} on {self.running_on}")
        print(f"Device's home directory: {self.home}")
        print(f"Project file structure:")
        print(f" - Root ........ {self.project_root} \n - Data Dir .... {self.data} \n - Notebooks ... {self.nbs}")
    
    def read_config(self) -> ConfigParser:
        """Read config from the configuration file if it exists and return an empty config if does not"""
        cfg = ConfigParser()
        if self.p2config.is_file(): 
            cfg.read(self.p2config)
        else:
            cfg.add_section('Infra')
        return cfg

    def register_as_local(self) -> ConfigParser:
        """Update the configuration file to register the machine as local machine"""
        cfg = self.read_config()
        os.makedirs(self.home/self._config_dir, exist_ok=True)
        cfg['Infra']['registered_as_local'] = 'True'
        with open(self.p2config, 'w') as fp:
            cfg.write(fp)
        return cfg

    def set_project_root(
        self, 
        p2project: str|Path,      # string or Path to the project directory. Can be absolute or relative to home
        data_dir: str = 'data'    # Directory name for data under project root
        ) -> ConfigParser:
        """Update the configuration file to set the project root"""
        # Build and validate the path to the project root
        if isinstance(p2project, str): 
            p2project = Path(p2project)
            if not p2project.is_absolute():
                p2project = self.home / p2project
        if not p2project.is_dir(): raise FileNotFoundError(f"{p2project.absolute()} does not exist")
        
        # Update the configuration file        
        cfg = self.read_config()
        os.makedirs(self.home/self._config_dir, exist_ok=True)
        cfg['Infra']['project_root'] = str(p2project.absolute())
        cfg['Infra']['data_dir'] = str(data_dir)
        with open(self.p2config, 'w') as fp:
            cfg.write(fp)
        self._project_root = p2project
        self._data = self.project_root / data_dir
        print(f"Project Root set to:   {self._project_root.absolute()}")
        print(f"Data directory set to: {self._data.absolute()}")
        return cfg

    @property
    def os(self) -> str: return sys.platform

    @property
    def project_root(self) -> Path:
        if self.is_local:
            return self._project_root
        elif self.is_colab:
            return self.gdrive / self._shared_project_dir
        elif os.getenv("GITHUB_ACTIONS") == "true":
            return self._project_root
        elif self.is_kaggle:
            raise NotImplementedError(f"ProjectFileSystem is not implemented for Kaggle yet")
        else:
            msg = """
                Not running locally, on Colab or on Kaggle. If running locally,
                register the machine as local by running `ProjectFileSystem().register_as_local()`
                """
            warnings.warn(msg)
            # raise ValueError('Not running locally, on Colab or on Kaggle')
            return self._project_root

    @property
    def data(self) -> Path: return self._data

    @data.setter
    def data(self, folder_name) -> None: self._data = self.project_root / folder_name
    
    @property
    def nbs(self) -> Path: return self.project_root / 'nbs'        

    @nbs.setter
    def nbs(self,folder_name) -> Path: return self.project_root / folder_name

    @property
    def p2config(self) -> Path: 
        # return self.home / self._config_dir / self._config_fname
        return self._p2config
        
    @property
    def is_local(self) -> bool:
        """Return `True` if the current machine was registered as a local machine"""
        cfg = self.read_config()
        return cfg['Infra'].getboolean('registered_as_local', False)

    @property
    def running_on(self) -> str:
        """Return the device on which this is run: local, colab, kaggle, ..."""
        if self.is_local: device = 'local computer'
        elif self.is_colab: device = 'colab'
        elif self.is_kaggle: device = 'kaggle'
        else: device = 'unknown cloud server'
        return device
    
    def readme(
        self, 
        dir_path:Path|None=None, # Path to the directory to inquire. If None, display readme file from project_root.
        ) -> None:
        """Display `readme.md` file or any other `.md` file in `dir_path`. 

        This provides a convenient way to get information on each direcotry content
        """
        if dir_path is None: 
            path = self.data
        elif validate_path(dir_path, path_type='dir'):
            path = dir_path
        else:
            raise ValueError(f"'dir_path' is not a directory: {dir_path.absolute()}")

        
        if path.is_relative_to(self.project_root):
            path_to_display = path.relative_to(self.project_root)
        else:
            path_to_display = path.absolute()
        display(HTML('<hr>'))
        # display(Markdown(f"ReadMe file for directory `{path.relative_to(self.project_root)}`:"))
        display(Markdown(f"ReadMe file for directory `{path_to_display}`:"))
        mdfiles = {p.stem: p for p in path.glob('*.md')}
        if mdfiles:
            mdfile = mdfiles.get('readme', None)
            if mdfile is None:
                mdfile = mdfiles.get(list(mdfiles.keys())[0])
            display(HTML('<hr>'))
            display(Markdown(filename=mdfile))
            display(HTML('<hr>'))
        else:
            print('No markdown file in this folder')
        

# %% ../nbs-dev/00_core.ipynb 36
class SqliteDatabase:
    """Manage a SQLite db file, execute SQL queries, return results, provide context manager functionality.

    Example usage as a context manager
    
    ```python
    db_path = Path('your_database.db')
    db = SqliteDb(db_path)

    with db as database:
        result = database.get_result("SELECT * FROM your_table")
        print(result)
    ```
    """

    def __init__(self, p2db:Path) -> None:
        self.p2db = p2db
        self.conn = None

    def connect(self) -> Connection:
        """Connect to the SQLite database"""
        self.conn = sqlite3.connect(self.p2db)
        return self.conn

    def execute(self, sql:str) -> Cursor:
        """Execute an SQL query and return the cursor after execution"""
        if self.conn is None: self.conn = self.connect()
        cursor = self.conn.cursor()
        cursor.execute(sql)
        return cursor

    def get_result(self, sql:str) -> list[Any]:
        """Execute an SQL query and return the result"""
        cursor = self.execute(sql)
        result = cursor.fetchall()
        return result

    def get_dataframe(self, sql:str) -> pd.DataFrame:
        """Wraps pandas.read_sql_query"""
        if self.conn is None: self.connect()
        df = pd.read_sql_query(sql, self.conn)
        return df
    
    def dataframe_to_table(
        self, 
        df:pd.DataFrame,         # DataFrame to write to the database
        table_name:str,          # Name of the database table to write to
        if_exists:Literal['fail','replace','append']='append',  # One of 'fail', 'replace', 'append'
        index:bool=False         # If True, write the DataFrame index as a column
        )-> None:
        """Wraps pandas.DataFrame.to_sql"""
        if self.conn is None: self.connect()
        df.to_sql(table_name, self.conn, if_exists=if_exists, index=index)
            
    def close(self) -> None:
        """Close the connection to the SQLite database"""
        if self.conn is not None:
            self.conn.close()
            self.conn = None

    def __enter__(self):
        """Enter the runtime context related to this object."""
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        """Exit the runtime context related to this object."""
        self.close()

    def list_columns(self,
                     name: str # name of a table or a view
                    ) -> list[str]:
        """Returns the list of columns in the table or view `name`"""
        query = f"PRAGMA table_info({name})"
        cursor = self.execute(query)
        cols = [row[1] for row in cursor.fetchall()]
        return cols

    def list_indexes(self) -> None:
        """List indexes in the database and the indexed columns"""
        print(f"List of indexes in database '{self.p2db.name}' and indexed columns:")
        if self.conn is None: self.connect()

        query = f"SELECT name, tbl_name FROM sqlite_master WHERE type='index' ;"
        indexes = self.get_result(sql=query)

        for index in indexes:
            print(f"- {index[0]} for table '{index[1]}':")
            cursor = self.execute(f"PRAGMA index_info({index[0]}) ;")
            for col_info in cursor.fetchall():
                print(f"    - {col_info[2]}")

    def print_schema(self) -> None:
        # 'type', 'name', 'tbl_name', 'rootpage', 'sql'
        # type can be 'table', 'view', 'index' or 'trigger'
        query = """
        SELECT type, name
        FROM sqlite_master 
        WHERE type IN ('table', 'view', 'trigger') 
        """
        cursor = self.execute(query)
        for t, name in cursor.fetchall():
            print(f"{name} ({t})")
            if t == 'table':
                print(' columns:',', '.join(self.list_columns(name)))
                indexes = self.get_result(f"PRAGMA index_list({name})")
                for index in indexes:
                    index_name = index[1]
                    print(f" index: {index_name}")
                    print('   indexed columns:',', '.join([row[2] for row in self.get_result(f"PRAGMA index_info({index_name})")]))
            if t =='view':
                print(' columns:', ','.join(self.list_columns(name)))
            print()


# %% ../nbs-dev/00_core.ipynb 41
class JsonDict(dict):
    """Dictionary whose current value is mirrored in a json file and can be initated from a json file
    
    `JsonDict` requires a path to json file at creation. An optional dict can be passed as argument.

    Behavior at creation:
    
    - `JsonDict(p2json, dict)` will create a `JsonDict` with key-values from `dict`, and mirrored in `p2json`
    - `JsonDict(p2json)` will create a `JsonDict` with empty dictionary and load json content if file exists

    Once created, `JsonDict` instances behave exactly as a dictionary
    """
    def __init__(
        self, 
        p2json: str|Path,               # path to the json file to mirror with the dictionary 
        dictionary: dict|None = None    # optional dictionary to initialize the JsonDict
        ):
        """Create dict from a passed dict or from json. Create the json file if required"""
        self.p2json = Path(p2json) if isinstance(p2json, str) else p2json
        if dictionary is None:
            if self.p2json.is_file():
                dictionary = self.load()
                self.initial_dict_from_json = True
            else:
                dictionary = dict()
                self.initial_dict_from_json = False
        super().__init__(dictionary)
        self.save()
    
    def __setitem__(self, __k:Any, v:Any) -> None:
        super().__setitem__(__k, v)
        self.save()

    def __delitem__(self, k:Any):
        super().__delitem__(k)
        self.save()

    def __repr__(self):
        txt1 = f"dict mirrored in {self.p2json.absolute()}\n"
        txt2 = super().__repr__()
        return txt1 + txt2

    def load(self)->dict:
        with open(self.p2json, 'r') as fp:
            return json.load(fp)

    def save(self):
        with open(self.p2json, 'w') as fp:
            json.dump(self, fp, indent=4)



# %% ../nbs-dev/00_core.ipynb 52
class JsonFileReader:
    """Mirror a JSON file and a dictionary"""
    def __init__(self, 
                 path:str|Path # path to the json file
                ):
        self.path = safe_path(path)
        with open(path, 'r') as fp:
            self.d = json.load(fp)
    
    def add_item(self, 
                 key:str,  # key for the new item
                 item:dict # new item to add to the json as a dict
                ):
        self.d[key] = item
        return self.d

    def save_to_file(self, path=None):
        if path is None: 
            path = self.path
        else:
            path = safe_path(path)

        with open(path, 'w') as fp:
            json.dump(self.d, fp, indent=4)

# %% ../nbs-dev/00_core.ipynb 62
def list_available_devices():
    # Check if CUDA is available
    cuda_available = torch.cuda.is_available()
    print(f"CUDA available: {cuda_available}")

    # List all available CUDA devices
    if cuda_available:
        num_cuda_devices = torch.cuda.device_count()
        print(f"Number of CUDA devices: {num_cuda_devices}")
        for i in range(num_cuda_devices):
            print(f"CUDA Device {i}: {torch.cuda.get_device_name(i)}")

    # Check for CPU availability
    cpu_available = torch.device('cpu')
    print(f"CPU available: {cpu_available}")

# %% ../nbs-dev/00_core.ipynb 66
class TextFileBaseReader:
    """Iterator going through a text file by chunks of `nlines` lines. Iterator can be reset to file start.
    
    The class is mainly intented to be extended, as it is for handling sequence files of various formats such as `FastaFileReader`.
    """

    def __init__(
        self,
        path: str|Path,  # path to the file
        nlines: int=1,   # number of lines on one chunk
    ):
        self.path = safe_path(path)
        self.nlines = nlines
        self.fp = None
        self.reset_iterator()
        
        # Attributes related to metadata parsing
        # Currently assumes the iterator generates a dictionary with key/values
        # TODO: extend to iterator output as simple string.
        self.text_to_parse_key = None
        project_root = ProjectFileSystem().project_root
        self.parsing_rules_json = project_root / 'default_parsing_rules.json'
        # self.parsing_rules_json = Path(f"{PACKAGE_ROOT}/default_parsing_rules.json")
        self.re_rule_name = None
        self.re_pattern = None       # regex pattern to use to parse text
        self.re_keys = None          # keys (re group names) to parse text

    def reset_iterator(self) -> None:
        """Reset the iterator to point to the first line in the file."""
        if self.fp is not None:
            self.fp.close()
        self.fp = open(self.path, 'r')
        self._chunk_nb = 0
        
    def __iter__(self):
        return self

    def _safe_readline(self) -> str:
        """Read a new line and handle end of file tasks."""
        if self.fp is None: raise RuntimeError(f"File {self.path} is not opened")
        line = self.fp.readline()
        if line == '':
            self.fp.close()
            raise StopIteration()
        return line

    def __next__(self) -> str:
        """Return one chunk of `nlines` text lines at the time"""
        lines = []
        for i in range(self.nlines):
            lines.append(self._safe_readline())
        self._chunk_nb = self._chunk_nb + 1
        return ''.join(lines)
    
    def print_first_chunks(
        self, 
        nchunks:int=3,  # number of chunks to print
    ) -> None:
        """Print the first `nchunk` chunks of text from the file.

        After printing, the iterator is reset again to its first line.
        """
        self.reset_iterator()
        for i, chunk in enumerate(self.__iter__()):
            if i > nchunks-1: break
            print(f"{self.nlines}-line chunk {i+1}")
            print(chunk)
        self.reset_iterator()
            
    def _parse_text_fn(
        self,
        txt:str,         # text to parse
        pattern:str,     # regex pattern to apply to parse the text
        keys:list[str],  # list of keys: keys are both the regex match group names and the corresponding output dict keys
        # TODO: remove the keys parameter and use the keys from the parsing rules json file
    )-> dict:            # parsed metadata in key/value format
        """Basic parser function parsing metadata from string, using regex pattern. Return a metadata dictionary."""

        #  TODO: refactor this function using groupdict() and using keys from the regex pattern. See code in comment below
        # p = re.compile(pattern)
        # keys = p.groupindex.keys()
        # match = p.search(txt)
        # if match is not None:
        #     metadata = match.groupdict() 

        matches = re.match(pattern, txt)
        metadata = {}
        if matches is not None:
            for g in sorted(keys):
                m = matches.group(g)
                metadata[g] = m.replace('\t', ' ').strip() if m is not None else None
        
        else: 
            # TODO: review hack below, to avoid error when missing metadata such as 'species name'.
            # Current code tries to recover by saving the entire line in the fist key, expected to be the seqid or refseid
            # if txt:
            #     metadata[keys[0]] = txt.replace('\t', ' ').strip() if txt is not None else None
            # else:
            raise ValueError(f"No match on this line")
        return metadata

    def parse_text(
        self,
        txt:str,                    # text to parse
        pattern:str|None=None,      # If None, uses standard regex pattern to extract metadata, otherwise, uses passed regex
        keys:list[str]|None=None,   # If None, uses standard regex list of keys, otherwise, uses passed list of keys (str)
    )-> dict:                       # parsed metadata in key/value format
        """Parse text using regex pattern and keys. Return a metadata dictionary.
        
        The passed text is parsed using the regex pattern. The method return a dictionary in the format:

            {
                'key_1': 'metadata 1',
                'key_2': 'metadata 2',
                ...
            }
        
        """
        # TODO: remove parameter keys as in _parse_text_fn
        if pattern is None and keys is None:
            if self.re_pattern is not None and self.re_keys is not None:
                return self._parse_text_fn(txt, self.re_pattern, self.re_keys)
            else:
                raise ValueError('attribute re_pattern and re_keys are still None')
        elif pattern is None or keys is None:
            raise ValueError('pattern and keys must be either both None or both have a value')
        else:
            return self._parse_text_fn(txt, pattern, keys)
        
        
    def set_parsing_rules(
        self,
        pattern: str|None=None,     # regex pattern to apply to parse the text, search in parsing rules json if None
        keys: list[str]|None=None,  # list of keys/group for regex, search in parsing rules json if None
        verbose: bool=False         # when True, provides information on each rule
    )-> None:
        """Set the standard regex parsing rule for the file.
        
        Rules can be set:
        
        1. manually by passing specific custom values for `pattern` and `keys`
        2. automatically, by testing all parsing rules saved in `parsing_rule.json` 
        
        Automatic selection of parsing rules works by testing each rule saved in `parsing_rule.json` on the first 
        definition line of the fasta file, and selecting the one rule that generates the most metadata matches.
        
        Rules consists of two parameters:
        
        - The regex pattern including one `group` for each metadata item, e.g `(?P<group_name>regex_code)`
        - The list of keys, i.e. the list with the name of each regex groups, used as key in the metadata dictionary
        
        This method updates the three following class attributes: `re_rule_name`, `re_pattern`, `re_keys`
      
        """
        # get the first definition line in the file to test the pattern
        # in base class, text_to_parse_key is not defined and automatic rule selection cannot be used
        # this must be handled in children classes
        if self.text_to_parse_key is None:
            msg = """
            `text_to_parse_key` is not defined in this class. 
            It is not possible to set a parsing rule.
            """
            warnings.warn(msg, category=UserWarning)
            return

        self.reset_iterator()
        first_output = next(self)
        text_to_parse = first_output[self.text_to_parse_key]
        divider_line = f"{'-'*80}"

        if pattern is not None and keys is not None:
            try:
                metadata_dict = self.parse_text(text_to_parse, pattern, keys)
                self.re_rule_name = 'Custom Rule'
                self.re_pattern = pattern
                self.re_keys = keys
                if verbose:
                    print(divider_line)
                    print(f"Custom rule was set for this instance.")
            except Exception as err: 
                raise ValueError(f"The pattern generates the following error:\n{err}")
                
        else:
            # Load all existing rules from json file
            with open(self.parsing_rules_json, 'r') as fp:
                parsing_rules = json.load(fp)
                
            # test all existing rules and keep the one with highest number of matches
            max_nbr_matches = 0
            for k, v in parsing_rules.items():
                re_pattern = v['pattern']
                re_keys = v['keys'].split(' ')
                try:
                    metadata_dict = self.parse_text(text_to_parse, re_pattern, re_keys)
                    nbr_matches = len(metadata_dict)
                    if verbose:
                        print(divider_line)
                        print(f"Rule <{k}> generated {nbr_matches:,d} matches")
                        print(divider_line)
                        print(re_pattern)
                        print(re_keys)

                    if len(metadata_dict) > max_nbr_matches:
                        self.re_pattern = re_pattern
                        self.re_keys = re_keys
                        self.re_rule_name = k    
                except Exception as err:
                    if verbose:
                        print(divider_line)
                        print(f"Rule <{k}> generated an error")
                        print(err)
                    else:
                        pass
            if self.re_rule_name is None:
                msg = """
        None of the saved parsing rules were able to extract metadata from the first line in this file.
        You must set a custom rule (pattern + keys) before parsing text, by using:
            `self.set_parsing_rules(custom_pattern, custom_list_of_keys)`
                """
                warnings.warn(msg, category=UserWarning)
            
            if verbose:
                print(divider_line)
                print(f"Selected rule with most matches: {self.re_rule_name}")

            # We used the iterator, now we need to reset it to make all lines available
            self.reset_iterator()


# %% ../nbs-dev/00_core.ipynb 87
class TextFileBaseIterator:
    """`TextFileBaseIterator` is a deprecated class, to be replaced by `TextFileBaseReader`"""
    def __init__(self, *args, **kwargs):
        msg = """
        `TextFileBaseIterator` is deprecated. 
        Use `TextFileBaseReader` instead, with same capabilities and more."""
        raise DeprecationWarning(msg)
