# -*- coding: utf-8 -*-
"""
Created on Mon Apr 17 17:15:36 2023

@author: sipost1
"""
import os
import glob
import csv
import argparse
from collections import defaultdict
from typing import Dict, TextIO

parser = argparse.ArgumentParser()
parser.add_argument("--inputfolder", "-i", type=str, required=True,
                    help="Input folder where report files can be found.")
parser.add_argument("--level", "-l", type=str, required=True,
                    help="Taxonomic level to extract from kraken2 report.")
parser.add_argument("--extension", "-e", default=".k2report",
                    help="Extension of kraken2 report files. Default: .k2report")
parser.add_argument("--outdir", "-o", type=str, default="./",
                    help="Output directory to save OTU table files. Default: current directory")
args = parser.parse_args()


def extract(file: str) -> Dict:
    """
    Requires a valid kraken2 report file.
    Returns a dictionary of taxas and readcounts.
    """

    taxons = defaultdict(dict)

    # Standard kraken2 output has 6 fields.
    # See kraken2 manual.
    indexes = {6: [3, 5], 8: [5, 7]}

    file_path = os.path.abspath(file)

    with open(file_path, "r") as ori:
        lines = ori.readlines()

    assert len(lines) > 0, "Empty file!"

    taxon_index, name_index = indexes[len(lines[0].split("\t"))]
    print(f"Reading {file}")
    print(f"Used indexes: field {taxon_index + 1} for taxon rank,\
 field {name_index + 1} for taxon name.")

    for line in lines:
        line_params = line.rstrip("\n").split("\t")

        read_count = line_params[1]
        taxon = line_params[taxon_index].strip()
        name = line_params[name_index].strip()

        taxons[taxon][name] = read_count

    return taxons


def read_in_files(directory: str, extension=args.extension) -> Dict:
    """
    Requires a directory where k2report files are.
    Returns a 3-depth dictionary.
    This function reads the input files into one dictionary.
    Finds all files within folder which end in given extension.
    """

    file_dictionary = defaultdict()
    report_files = glob.glob(f"{directory}/*{extension}")
    assert len(
        report_files) > 0, "No report file found!\nPlease check filename extension and input dir!"

    for file in report_files:
        file = os.path.abspath(file)
        file_dictionary[os.path.basename(file).rstrip(
            f"{extension}")] = extract(file)

    return file_dictionary


def create_otu_table(level: str, file_sample_dict: Dict, outdir="./") -> TextIO:
    """

    Parameters
    ----------
    level : str
        One character level (O1,G,S).
    file_sample_dict : Dict
        Dictionary generated by read_in_files.
    outdir : str, optional
        Directory to save the OTU table files, by default "./"

    Returns
    -------
    TextIO
        OTU table for given level.

    """
    rearranged_dict = defaultdict(dict)

    level = level.upper()
    sample_taxa = {sample: level_list[level]
                   for sample, level_list in file_sample_dict.items()}

    for sample, taxon_list in sample_taxa.items():
        for taxon in taxon_list:
            rearranged_dict[taxon][sample] = taxon_list[taxon]

    headers = ["otu"] + list(sample_taxa.keys())

    outfile_name = f"otu_table_{level}.csv"
    with open(f"{os.path.join(outdir, outfile_name)}", "w", newline="") as csv_file:
        writer = csv.writer(csv_file)

        # Write the headers to the CSV file
        writer.writerow(headers)

        # Write the values from the dictionary to the CSV file
        for otu, inner_dict in rearranged_dict.items():
            # Initialize a list for this row
            row = [otu]

            # Add the values to the row in the order specified by the headers
            for header in headers[1:]:
                if header in inner_dict:
                    row.append(inner_dict[header])
                else:
                    row.append(0)

            # Write the completed row to the CSV file
            writer.writerow(row)


if __name__ == "__main__":
    file_dict = read_in_files(args.inputfolder)
    create_otu_table(args.level, file_dict, outdir=args.outdir)
    print("Done!")
