from typing import Dict
from pathlib import Path
from typing import Optional
import csv


def write_results(
    result_dict: Dict[str, Dict[str, str]],
    trans_dict: Dict[str, str],
    input_filepath: Path,
    output_filepath: Path,
    write_csv: Optional[bool] = False,
) -> Dict[str, Dict[str, str]]:

    """
    Returns processed output by combining results_dict (predicted values) and trans_dict (compound IDs).
    Optionally writes results to a CSV file.
    """

    # obtaining all possible column names
    acd_columns = []
    counter = 0
    for key, value in result_dict.items():
        for col, value1 in value.items():
            if col not in acd_columns:
                acd_columns.append(col)
        counter += 1
        if counter == 10 ** 4:
            break

    # filling in missing columns
    for key, value in result_dict.items():
        for col in acd_columns:
            if col not in value:
                result_dict[key][col] = "NaN"

    # translating ID back to original IDs as provided in input file
    trans_result_dict = {}
    for cp_id, props in result_dict.items():
        trans_result_dict[trans_dict[cp_id]] = props

    # writting to csv
    if write_csv is True:
        acd_columns.append("compound_id")
        with open(output_filepath, "w") as f:
            w = csv.DictWriter(f, acd_columns)
            w.writeheader()
            for k in trans_result_dict:
                w.writerow(
                    {col: trans_result_dict[k].get(col) or k for col in acd_columns}
                )

    return trans_result_dict