#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
input: file names of MPA's metaprotein report, style yaml and output
output: protein group yaml for prophane
"""

import pandas as pd

from utils.search_result_parsing.helper_classes import ProteomicSearchResultReaderBaseClass, InputFile


class MpaPortableTableReader(ProteomicSearchResultReaderBaseClass):
    mandatory_style_fields = ProteomicSearchResultReaderBaseClass.mandatory_style_fields[:]
    mandatory_style_fields.remove("sample_cols")
    mandatory_style_fields.append("sample_col")

    def read_file_into_dataframe(self):
        file_obj = InputFile(self.result_table)
        required_col_names = [self.style_dict['proteins_col'],
                              self.style_dict['quant_col']]
        optional_col_names = [self.style_dict['sample_col'], ]
        raw_df = pd.read_csv(self.result_table, sep=self.style_dict['field_sep'], encoding=file_obj.encoding)
        if self.style_dict['sample_col'] in raw_df.columns:
            df = raw_df[required_col_names + optional_col_names].copy()
            df['sample'] = df[self.style_dict['sample_col']]
        else:
            print(f"No sample column found (looked for '{self.style_dict['sample_col']}').\n"
                  "Assigning all identifications to 'sample1'.")
            df = raw_df[required_col_names].copy()
            df['sample'] = "sample1"

        df['proteins'] = df[self.style_dict['proteins_col']].apply(lambda x: x.split(self.style_dict['proteins_sep']))

        df_standard = pd.DataFrame()
        df_standard['proteins'] = df['proteins']
        df_standard['sample'] = df['sample']
        df_standard['quant'] = df[self.style_dict['quant_col']]
        return df_standard
