from __future__ import annotations
from scikit_fluorescence import typing

import numpy as np

#  N labeling sites
#  experiment -> N potential labeling sites
#  experiment -> data (fluorescence decay, cross linking, pda)
#  experiment(labeling site distribution) -> score
#  experiments(labeling site distribution) -> total score = sum of scores


class ExperimentalOutcome(object):

    def __mul__(self, a):
        return ExperimentalOutcome(self)

    def __add__(self, other):
        return ExperimentalOutcome(self)


class LabelingSite(object):

    id: int
    coordinates: np.ndarray
    residue_seq_number: int
    chain_id: str
    atom_name: str


class LabeledSite(object):

    id: int
    labeling_site: LabeledSite
    dye_parameters: typing.Dict


class LabelSitePair(object):

    label_site_1: LabeledSite
    label_site_2: LabeledSite
    probability_matrix: np.ndarray  # ?
    # DA, D0, 0D, AD, AA, A0, 0A

    def expected_experimental_outcome(self) -> ExperimentalOutcome:
        """Forward model
        """
        #re = compute_experimental_outcome(self.label_site_1, self.label_site_2)
        re = ExperimentalOutcome()
        return re


class Experiment(object):

    label_site_pair_fractions: typing.List[
        typing.Tuple[float, LabelSitePair]
    ]

    def expected_experimental_outcome(self):
        """Forward model
        """
        re = ExperimentalOutcome()
        for fraction, label_site_pair in self.label_site_pair_fractions:
            re += label_site_pair.expected_experimental_outcome() * fraction
        return re


# compute distances between labeling sites - compare against experiment - score

# update positions of labeled residues
# iterate through experimental data and compute experiments that are needed to score

