# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03_seq_data_manager.ipynb.

# %% auto 0
__all__ = ['WgsData']

# %% ../nbs/03_seq_data_manager.ipynb 3
# standard libs
import os
import sys
import re

# Common to template
# add into settings.ini, requirements, package name is python-dotenv, for conda build ensure `conda config --add channels conda-forge`
import dotenv  # for loading config from .env files, https://pypi.org/project/python-dotenv/
import envyaml  # Allows to loads env vars into a yaml file, https://github.com/thesimj/envyaml
import fastcore  # To add functionality related to nbdev development, https://github.com/fastai/fastcore/
from fastcore import (
    test,
)
from fastcore.script import (
    call_parse,
)  # for @call_parse, https://fastcore.fast.ai/script
import json  # for nicely printing json and yaml

# import functions from core module (optional, but most likely needed).
from . import core

# Project specific libraries
from pathlib import Path
import pandas
import numpy

# %% ../nbs/03_seq_data_manager.ipynb 7
class WgsData:

    def __init__(self, file_paths: dict[dict], sample_data: dict[dict] = None):
        """
        Default initialization is from file paths provided in a dictionary with sample name as keys, like this:
        file_paths =
        {<sample_name>: {
            "assembly_file": <assembly_file_path>,
            "r1_file": <r1_file_path>,
            "r2_file": <r2_file_path>
            }
        }
        and metadata in a similar format with sample names as keys and a dictionary of attribute: value pairs as values.
        {<sample_name>: {
            "serotype": "a",
            "MLST": "15",
            "source": blood,
            ....
            }
        }



        To initialize from folders without using a metadata sheet, use WgsData.from_folders(<assembly_data_folder>, <read_data_folder>)
        Both assembly_data_folder and read_data_folder can be set to None if you have only reads or only assemblies.
        """
        self.sample_data = sample_data
        self.sample_names = file_paths.keys()
        self.file_paths = file_paths
        self.paired_end_read_files = []
        self.assembly_files = []
        for files in file_paths.values():
            if "r1_file" in files and "r2_file" in files:
                self.paired_end_read_files.append((files["r1_file"], files["r2_file"]))
            else:
                self.paired_end_read_files.append((None, None))
            if "assembly_file" in files:
                self.assembly_files.append(files["assembly_file"])
            else:
                self.assembly_files.append(None)

        if sample_data:
            self.sample_data_df = pandas.DataFrame.from_dict(
                sample_data, orient="index"
            )
            self.sample_data_df.index = self.sample_names
            self.sample_data_columns = list(self.sample_data_df.columns.values)
        else:
            self.sample_data_columns = None
            self.sample_data_df = None

    @classmethod
    def from_samplesheet(cls, sample_sheet: Path, base_folder: Path = None):
        """
        Load data from samplesheet into dictionary.

        """

        df = pandas.read_csv(sample_sheet, sep="\t")
        file_paths_df = df[["sample_name", "assembly_file", "r1_file", "r2_file"]]
        file_paths = {}
        sample_name_count = {}
        uniq_sample_names = []
        for index, row in file_paths_df.iterrows():
            sample_name = row["sample_name"]
            if not sample_name in sample_name_count:
                sample_name_count[sample_name] = 1
                file_paths[sample_name] = row.to_dict()
                uniq_sample_names.append(sample_name)
            else:
                sample_name_count[sample_name] += 1
                file_paths[f"{sample_name}_{sample_name_count[sample_name]}"] = (
                    row.to_dict()
                )
                uniq_sample_names.append(
                    f"{sample_name}_{sample_name_count[sample_name]}"
                )
                print(
                    f"Duplicate sample name {sample_name} in provided sample sheet. Using sample name {sample_name}_{sample_name_count[sample_name]} for uniquification"
                )
        if base_folder:
            base_folder = Path(base_folder).absolute()
            for files in file_paths.values():
                files["assembly_file"] = base_folder.joinpath(files["assembly_file"])
                files["r1_file"] = base_folder.joinpath(files["r1_file"])
                files["r2_file"] = base_folder.joinpath(files["r2_file"])
        else:
            base_folder = Path(sample_sheet).parent
            for sample, files in file_paths.items():
                if not Path(files["assembly_file"]).is_absolute():
                    files["assembly_file"] = base_folder.joinpath(
                        files["assembly_file"]
                    ).resolve()
                if not Path(files["r1_file"]).is_absolute():
                    files["r1_file"] = base_folder.joinpath(files["r1_file"]).resolve()
                if not Path(files["r2_file"]).is_absolute():
                    files["r2_file"] = base_folder.joinpath(files["r2_file"]).resolve()
        df["sample_name_uniq"] = uniq_sample_names
        df.set_index("sample_name_uniq", inplace=True)
        sample_data = df.to_dict("index")
        return cls(file_paths=file_paths, sample_data=sample_data)

    @classmethod
    def from_folders(
        cls,
        assembly_data_folder: Path = None,
        paired_end_read_data_folder: Path = None,
        fasta_file_pattern="(?P<sample_name>.+?)(\.fa|\.fna|\.fasta)$",
        paired_end_reads_pattern="(?P<sample_name>.+?)(?P<sample_number>(_S[0-9]+)?)(?P<lane>(_L[0-9]+)?)[\._]R?(?P<paired_read_number>[1|2])(?P<set_number>(_[0-9]+)?)(?P<file_extension>\.fastq\.gz)",
    ):
        if assembly_data_folder:
            assembly_data_dict = cls.parse_folder_for_fasta_files(
                data_folder=assembly_data_folder, fasta_file_pattern=fasta_file_pattern
            )
        if assembly_data_folder and len(assembly_data_dict) > 0:
            assembly_df = pandas.DataFrame.from_dict(assembly_data_dict, orient="index")
            assembly_df.columns = ["assembly_file"]
        else:
            assembly_df = pandas.DataFrame(
                {
                    "assembly_file": [],
                }
            )
        if paired_end_read_data_folder:
            paired_end_read_data_dict, unmatched_read_files = (
                cls.parse_folder_for_paired_end_reads(
                    data_folder=paired_end_read_data_folder,
                    paired_end_reads_pattern=paired_end_reads_pattern,
                )
            )

            paired_end_read_df = pandas.DataFrame(paired_end_read_data_dict).transpose()
        else:
            paired_end_read_df = pandas.DataFrame(
                {"sample_name": []}, index="assembly_file"
            )

        combined_df = assembly_df.merge(
            paired_end_read_df, how="outer", left_index=True, right_on="sample_name"
        )
        combined_df.replace(numpy.nan, None, inplace=True)
        combined_df = combined_df[
            ["sample_name", "assembly_file", "r1_file", "r2_file"]
        ]
        # print(combined_df)
        combined_dict = combined_df.to_dict("index")

        return cls(file_paths=combined_dict)

    def __iter__(self):
        if self.sample_data is None:
            for sample_name, files in self.file_paths.items():
                yield sample_name, files, None
        else:
            for sample_name, files in self.file_paths.items():
                yield sample_name, files, self.sample_data[sample_name]

    def get_missing_files(self):
        missing_assembly_files = {}
        missing_read_files = {}
        for sample_name, assembly_file, r1_file, r2_file in self.filepaths():
            if assembly_file is not None and not assembly_file.exists():
                missing_assembly_files[sample_name] = assembly_file
            if r1_file and r2_file and (not r1_file.exists() or not r2_file.exists()):
                missing_read_files[sample_name] = (r1_file, r2_file)
        return missing_assembly_files, missing_read_files

    def wgs_files_exist(self):
        missing_assembly_files, missing_read_files = self.get_missing_files()
        if len(missing_assembly_files) == 0 and len(missing_read_files) == 0:
            return True
        else:
            return False

    def from_dataframe(cls, pandas_dataframe: pandas.DataFrame):
        sample_data = pandas_dataframe.to_dict("index")
        sample_data.replace(numpy.nan, None, inplace=True)
        return cls(file_paths=sample_data)

    def filepaths(self):
        for i, sample_name in enumerate(self.sample_names):
            yield sample_name, self.assembly_files[i], self.paired_end_read_files[i][
                0
            ], self.paired_end_read_files[i][1]

    def readpaths(self):
        for i, sample_name in enumerate(self.sample_names):
            yield sample_name, self.paired_end_read_files[
                0
            ], self.paired_end_read_files[1]

    def assemblypaths(self):
        for i, sample_name in enumerate(self.sample_names):
            yield sample_name, self.assembly_files[i]

    def __repr__(self):
        if self.sample_data:
            return f"< WgsData object with {len(self.sample_names)} samples and {len(self.sample_data_columns)} metadata columns >"
        else:
            return f"< WgsData object with data from {len(self.sample_names)} samples >"

    @staticmethod
    def parse_folder_for_paired_end_reads(
        data_folder: Path,
        paired_end_reads_pattern: str = "(?P<sample_name>.+?)(?P<sample_number>(_S[0-9]+)?)(?P<lane>(_L[0-9]+)?)[\._]R?(?P<paired_read_number>[1|2])(?P<set_number>(_[0-9]+)?)(?P<file_extension>\.fastq\.gz)",
    ) -> tuple[dict[dict], set]:
        """
        Parses a provided folder for paired end reads matching the provided regex pattern
        Returns a dictionary with sample names as keys and a dict with keys "r1_file" and "r2_file" and filenames as values
        In case of duplicate sample names, the <sample_name>_<sample_number> combination will be used instead
        If <sample_name>_<sample_number> is also not unique, the key used for sample name will simply be the name of the r1_file
        """
        sample_list = []
        unmatched_files = set()
        sample_names = set()
        duplicate_sample_names = set()
        duplicate_sample_names_uniq = set()
        data_folder = Path(data_folder).resolve()
        for f in data_folder.iterdir():
            fname = f.name
            fname_match = re.match(paired_end_reads_pattern, fname)
            if fname_match:
                sample_name = fname_match.group("sample_name")
                if fname_match.group("paired_read_number") is not None:
                    read_number = fname_match.group("paired_read_number")
                    if fname_match.group("sample_number") is not None:
                        sample_name_uniq = (
                            f"{sample_name}{fname_match.group('sample_number')}"
                        )
                    else:
                        sample_name_uniq = f"{sample_name}"
                    if read_number == "1":
                        r1_file = f.absolute()
                        r2_file = data_folder.joinpath(
                            fname[: fname_match.start("paired_read_number")]
                            + "2"
                            + fname[fname_match.end("paired_read_number") :]
                        )
                        if r2_file.exists():
                            sample_list.append(
                                (sample_name, sample_name_uniq, r1_file, r2_file)
                            )
                            if sample_name in sample_names:
                                duplicate_sample_names.add(sample_name)
                                if sample_name_uniq in duplicate_sample_names:
                                    duplicate_sample_names_uniq.add(sample_name_uniq)
                            else:
                                sample_names.add(sample_name)
                        else:
                            unmatched_files.add(fname)
                    if read_number == "2":
                        r1_file = data_folder.joinpath(
                            fname[: fname_match.start("paired_read_number")]
                            + "1"
                            + fname[fname_match.end("paired_read_number") :]
                        )
                        if not r1_file.exists():
                            unmatched_files.add(fname)
        sample_file_dict = {}
        for sample_name, sample_name_uniq, r1_file, r2_file in sample_list:
            if sample_name in duplicate_sample_names:
                if sample_name in duplicate_sample_names_uniq:
                    sample_file_dict[r1_file] = {"r1_file": r1_file, "r2_file": r2_file}
                else:
                    sample_file_dict[sample_name_uniq] = {
                        "sample_name": sample_name,
                        "r1_file": r1_file,
                        "r2_file": r2_file,
                    }
            else:
                sample_file_dict[sample_name] = {
                    "sample_name": sample_name,
                    "r1_file": r1_file,
                    "r2_file": r2_file,
                }

        return sample_file_dict, unmatched_files

    @staticmethod
    def parse_folder_for_fasta_files(
        data_folder: Path,
        fasta_file_pattern: str = "(?P<sample_name>.+?)(\.fa|\.fna|\.fasta)$",
    ) -> tuple[dict[dict], set]:
        """
        Parses a provided folder for paired end reads matching the provided regex pattern
        Returns a dictionary with sample names as keys and a dict with keys "r1_file" and "r2_file" and filenames as values
        In case of duplicate sample names, the <sample_name>_<sample_number> combination will be used instead
        If <sample_name>_<sample_number> is also not unique, the key used for sample name will simply be the name of the r1_file
        """
        sample_file_dict = {}
        data_folder = Path(data_folder).resolve()
        for f in data_folder.iterdir():
            fname = f.name
            fname_match = re.match(fasta_file_pattern, fname)
            if fname_match:
                sample_name = fname_match.group("sample_name")
                sample_file_dict[sample_name] = f
        return sample_file_dict
