import pandas as pd
from Chamber_client import fetch_molecule_by_chembl_id, fetch_activity_by_target, search_molecules

class MultiSourceDataSource:
    """
    Advanced Multi-Source DataSource for Pharmalyzer.
    Supports ChEMBL and Papyrus datasets.
    Collects SMILES and target properties, ready for QSPRDataset.
    """

    def __init__(self, sources, properties, chembl_ids=None, target_names=None, papyrus_file=None):
        """
        Parameters:
        - sources: list of sources, e.g., ['ChEMBL', 'Papyrus']
        - properties: list of target properties, e.g., ['IC50','Ki']
        - chembl_ids: list of ChEMBL IDs (optional)
        - target_names: list of target names (optional)
        - papyrus_file: path to Papyrus CSV/TSV file (required if 'Papyrus' in sources)
        """
        self.sources = sources
        self.properties = properties
        self.chembl_ids = chembl_ids
        self.target_names = target_names
        self.papyrus_file = papyrus_file

    def fetch_data(self):
        all_data = []

        for source in self.sources:
            if source == 'ChEMBL':
                all_data.append(self._fetch_chembl())
            elif source == 'Papyrus':
                if not self.papyrus_file:
                    raise ValueError("Papyrus file path must be provided")
                all_data.append(self._fetch_papyrus())
            else:
                raise ValueError(f"Unsupported source: {source}")

        if all_data:
            return pd.concat(all_data, ignore_index=True)
        else:
            return pd.DataFrame()

    # -------- ChEMBL Fetching --------
    def _fetch_chembl(self):
        data = []
        ids = self.chembl_ids if self.chembl_ids else []
        targets = self.target_names if self.target_names else []

        # Fetch by IDs
        for cid in ids:
            mol = fetch_molecule_by_chembl_id(cid)
            if mol and 'smiles' in mol:
                row = {'ChEMBL_ID': cid, 'SMILES': mol['smiles']}
                for prop, target in zip(self.properties, targets):
                    activity = fetch_activity_by_target(cid, target) if target else None
                    row[prop] = activity
                data.append(row)

        # Fetch by target names
        for target in targets:
            molecules = search_molecules(target)
            for mol in molecules:
                cid = mol.get('chembl_id')
                smiles = mol.get('smiles')
                if smiles is not None:
                    row = {'ChEMBL_ID': cid, 'SMILES': smiles}
                    for prop in self.properties:
                        activity = fetch_activity_by_target(cid, target)
                        row[prop] = activity
                    data.append(row)

        return pd.DataFrame(data)

    # -------- Papyrus Fetching --------
    def _fetch_papyrus(self):
        df = pd.read_csv(self.papyrus_file)
        columns = ['SMILES'] + self.properties
        df_filtered = df[columns].dropna(subset=self.properties)
        df_filtered.reset_index(drop=True, inplace=True)
        return df_filtered

