# FTreesCSVHelper.py - Data handling routines developed by Optibrium for integration
# with StarDrop software
# Use of this code is subject to the terms of the StarDrop License Agreement and
# use, copying or redistribution of this file for use by other persons or
# organisations is prohibited
#
# Copyright(C) Optibrium Ltd 2026
# support@optibrium.com

"""
Module provides some functions to assist in parsing CSV type data into a
StarDrop data set.
"""
import app
import string
import csv


class ColumnInfo:
    """
    A small class to represent a column of information to be loaded into StarDrop.
    """

    def __init__(self, name, index):
        """
        Create a ColumnInfo object with a name and column index.

        :param name: The name of the column to appear in StarDrop.
        :param index: The column index of the columns data in the CVS data.
        """
        self.name = name
        self.column_index = index
        self.error_index = -1

    def __repr__(self):
        return "Name: %s, index %i, error %i" % (self.name, self.column_index, self.error_index)


def check_type(values, name, config):
    """
    Deduce the desired type from a column of data. Assumes data is continuous by
    default.

    :param values: A list of strings representing the data in the column.
    :param name: The name of the column.
    :param config: A configuration module that can be used to control the process.
    :return: A StarDrop data type (cont, mol, or string)
    """
    # check for molecules/structures
    rv = app.cont

    if 'name' in name:
        rv = app.string
    elif 'smiles' in name:
        rv = app.mol
    else:

        if name.upper().find("STRUCT") != -1 or name.upper().find("SMILES") != -1:
            rv = app.mol
            return rv
        # check for numerical values - check first TYPE_CHECK_DEPTH, and look for
        # at least TYPE_CHECK_FRACTION to be numbers
        count = 0.0
        numbers = 0.0
        index = 0
        while count < config.TYPE_CHECK_DEPTH:
            if index < len(values):
                # check for a null - count as valid number
                if values[index]:
                    count += 1.0
                    try:
                        _ = float(values[index])
                        numbers += 1.0
                    except ValueError:
                        if (values[index].lower() == 'invalid' or
                                # count this entry as representing a number column with invalid entries
                                len(values[index]) > 0 and values[index][0] in ('<', '>', '=')):
                            # count a EQ, GT or LT operator as being a number
                            numbers += 1.0
                index += 1
            else:
                break
        if count > 0:
            if numbers / count >= config.TYPE_CHECK_FRACTION:
                rv = app.cont
            else:
                rv = app.string
    return rv


def extract_result(result):
    """
    Extract a value and a qualifier from an input string.

    :param result: The input string.
    :return: A tuple of qualifier and value (as a float)
    """
    # if the data is empty or invalid at this point an exception will be thrown and caught in loadData
    qualifier = '\x00'
    main = result.rstrip(string.ascii_letters)
    if main[0] == '<':
        qualifier = '\x01'
        main = main.lstrip('<')
        main = main.lstrip()
    elif main[0] == '>':
        qualifier = '\x02'
        main = main.lstrip('>')
        main = main.lstrip()
    elif main[0] == '=':
        main = main.lstrip('=')
        main = main.lstrip()
    value = float(main)
    return qualifier, value


def extract_value_from_list(results):
    """
    Get a data entry from a list of values by calculating a mean.

    :param results: The values as strings.
    :return: A numeric data entry containing the average value, or an invalid entry.
    """
    count = 0
    total = 0.0
    for r in results:
        try:
            value = float(r.lstrip().rstrip())
            total += value
            count += 1
        except ValueError:
            pass
    if count > 0:
        return app.Dataset.cont(total / count, 0.0)
    else:
        return app.Dataset.createInvalidEntry(app.cont)


def create_cont_entry(entry, error_entry='0.0'):
    """
    Create a numeric entry from strings for the value and the error.

    :param entry: The entry value as a string , which may include a qualifier.
    :param error_entry: The error value as a string.
    :return: A numeric data entry for a StarDrop data set.
    """
    error = 0.0
    qualifier = '\x00'
    try:
        error = float(error_entry)
        value = float(entry)
    except ValueError:
        results = entry.split(',')
        if len(results) > 1:
            return extract_value_from_list(results)
        qualifier, value = extract_result(results[0].lstrip().rstrip())
    return 12 + 2 * 8 + 1, (True, [], []), (qualifier, value, error)


def create_data_set(output, config, id_column_name, base_url):
    """
    Take the content of a csv file or csv-like output from another source,
    and parse into a StarDrop data set.

    :param output: The csv data as a string.
    :param config: A configuration module that can be used to control the process.
    :param id_column_name: The name to be applied to the result name column in the data set.
    :param base_url: The URL to use when creating links.
    :return: A python representation of a StarDrop data set.
    """

    with open(output, 'r') as csvfile:
        csv_reader = csv.reader(csvfile, delimiter=',', quotechar='"')
        first = True
        headers = []
        table = []

        for row in csv_reader:
            if first:
                first = False
                i = 0
                for header in row:
                    table.append([])
                    headers.append(ColumnInfo(header, i))
                    i += 1
            else:
                if len(row) > len(headers):
                    raise ValueError('Incomplete data')
                elif len(row) < len(headers):
                    while len(row) < len(headers):
                        row.append('')

                i = 0
                for item in row:
                    table[i].append(item)
                    i += 1

        # now build the columns from the data in the table
        ds = app.Dataset()
        ds.cols = []
        for header in headers:
            if header.name.rstrip('"').lstrip('"') == '#result-smiles':
                header_corrected = 'Result structure'
            elif header.name.rstrip('"').lstrip('"') == 'result-rank':
                header_corrected = 'Result rank'
            elif header.name.rstrip('"').lstrip('"') == 'similarity':
                header_corrected = 'Similarity'
            elif header.name.rstrip('"').lstrip('"') == 'result-name':
                header_corrected = id_column_name
            elif header.name.rstrip('"').lstrip('"') == 'query-smiles':
                header_corrected = 'Query SMILES'
            elif header.name.rstrip('"').lstrip('"') == 'similarity-descriptor-smiles':
                header_corrected = 'similarity-descriptor-smiles'
            elif header.name.rstrip('"').lstrip('"') == 'similarity-descriptor':
                header_corrected = 'similarity-descriptor'
            elif header.name.rstrip('"').lstrip('"') == 'source':
                header_corrected = 'Source'
            elif header.name.rstrip('"').lstrip('"') == 'reaction-name':
                header_corrected = 'Reaction name'
            elif header.name.rstrip('"').lstrip('"') == 'reagent1-name':
                header_corrected = 'Reagent 1 name'
            elif header.name.rstrip('"').lstrip('"') == 'reagent1-smiles':
                header_corrected = 'Reagent 1 SMILES'
            elif header.name.rstrip('"').lstrip('"') == 'reagent2-name':
                header_corrected = 'Reagent 2 name'
            elif header.name.rstrip('"').lstrip('"') == 'reagent2-smiles':
                header_corrected = 'Reagent 2 SMILES'
            elif header.name.rstrip('"').lstrip('"') == 'reagent3-name':
                header_corrected = 'Reagent 3 name'
            elif header.name.rstrip('"').lstrip('"') == 'reagent3-smiles':
                header_corrected = 'Reagent 3 SMILES'
            elif header.name.rstrip('"').lstrip('"') == 'reagent4-name':
                header_corrected = 'Reagent 4 name'
            elif header.name.rstrip('"').lstrip('"') == 'reagent4-smiles':
                header_corrected = 'Reagent 4 SMILES'
            elif header.name.rstrip('"').lstrip('"') == 'reagent5-name':
                header_corrected = 'Reagent 5 name'
            elif header.name.rstrip('"').lstrip('"') == 'reagent5-smiles':
                header_corrected = 'Reagent 5 SMILES'
            else:
                header_corrected = header.name.rstrip('"').lstrip('"')
            col = []
            if header_corrected in ('similarity-descriptor', 'similarity-descriptor-smiles'):
                ds.columns.append(
                    app.Dataset.meta(header_corrected, app.string, ''))
            else:
                if len(table[:-1]) > 0:
                    column_type = app.cont
                    if header.error_index == -1:
                        column_type = check_type(table[header.column_index], header.name, config)
                    try:
                        if column_type == app.cont:
                            ds.columns.append(app.Dataset.createNumberHeader(header_corrected))
                        else:
                            ds.columns.append(
                                app.Dataset.meta(header_corrected, column_type))

                        if header.error_index == -1:
                            for entry in table[header.column_index]:

                                if entry is None:
                                    entry = ''
                                if column_type == app.mol:
                                    # TODO: use regex here?
                                    entry = entry.replace('[R1*]', '*').replace(
                                        '[R2*]', '*').replace(
                                        '[R3*]', '*').replace(
                                        '[R4*]', '*').replace(
                                        '[R5*]', '*')
                                    try:
                                        col.append(app.Dataset.molecule(entry))
                                    except Exception as e:
                                        # TODO: what is thrown here?
                                        print(e)
                                elif column_type == app.string:
                                    # TODO: make this bit general
                                    if header_corrected == id_column_name:
                                        if entry and "__" not in entry:
                                            entry = "<a href=\"{0}/{1}\">{1}</a>".format(base_url, entry)

                                    col.append(app.Dataset.string(entry))
                                else:
                                    try:
                                        col.append(create_cont_entry(entry))
                                    except Exception as e:
                                        col.append(app.Dataset.createInvalidEntry(app.cont))
                                        if str(e).find("Entry contains a list") != -1:
                                            raise
                        else:
                            for value, error in zip(table[header.column_index], table[header.error_index]):
                                try:
                                    col.append(create_cont_entry(value, error))
                                except ValueError:
                                    col.append(app.Dataset.createInvalidEntry(app.cont))
                    except Exception as e:
                        # add as a string column instead
                        print(e)
                        ds.columns.pop()
                        col = []
                        ds.columns.append(
                            app.Dataset.meta(header_corrected, app.string))
                        for entry in table[header.column_index]:
                            if entry is None:
                                entry = ''
                            col.append(app.Dataset.string(entry))
                else:
                    ds.columns.append(
                        app.Dataset.meta(header_corrected, app.string))
            ds.cols.append(col)

        list_of_cols_to_del = []

        # identify any empty columns, put them into a list
        for i in range(len(headers)):
            valid_entry = False
            if i < 4:
                valid_entry = True
            else:
                for field in ds.cols[i]:
                    if field[2] == '\x00':
                        continue
                    elif field[2][0] == "":
                        continue
                    else:
                        valid_entry = True
                        break

            if not valid_entry:
                list_of_cols_to_del.append(i)

        # reverse the list of empty columns to delete them without affecting indexing
        reverse_list = list_of_cols_to_del[::-1]

        # delete the empty columns
        if reverse_list:
            for col_num in reverse_list:
                ds.delColumn(col_num)

        return ds
