#!/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 MpaPortableMultisampleTableReader(ProteomicSearchResultReaderBaseClass):
    mandatory_style_fields = ProteomicSearchResultReaderBaseClass.mandatory_style_fields[:]
    mandatory_style_fields.append("sample_col_tag")
    mandatory_style_fields.remove("sample_cols")
    mandatory_style_fields.remove('quant_col')

    def read_file_into_dataframe(self):
        file_obj = InputFile(self.result_table)
        required_col_names = [self.style_dict['proteins_col'], ]

        raw_df = pd.read_csv(self.result_table, sep=self.style_dict['field_sep'], encoding=file_obj.encoding)

        sample_col_names = [col for col in raw_df.columns if col.startswith(self.style_dict['sample_col_tag'])]
        if not sample_col_names:
            raise ValueError(f"Found no sample columns that start with '{self.style_dict['sample_col_tag']}'"
                             f"\tin file {self.result_table}"
                             f"\tMake sure you used MPA-server version 3.4 or higher")
        required_col_names += sample_col_names

        df = raw_df[required_col_names].copy()
        del raw_df
        sample_renaming_dict = {old: old[len(self.style_dict['sample_col_tag']):] for old in sample_col_names}
        df.rename(columns=sample_renaming_dict, inplace=True)
        df = pd.melt(df, id_vars=[self.style_dict['proteins_col']], value_vars=sample_renaming_dict.values(),
                     var_name='sample', value_name='quant')
        df.rename(columns={self.style_dict['proteins_col']: "proteins"}, inplace=True)

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

        return df
