#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

# pyre-strict

from enum import Enum
from unittest import mock

import numpy as np
import pandas as pd
from ax.core.arm import Arm
from ax.core.generator_run import GeneratorRun
from ax.metrics.chemistry import ChemistryMetric, ChemistryProblemType
from ax.utils.common.testutils import TestCase
from ax.utils.testing.core_stubs import get_trial


class DummyEnum(Enum):
    DUMMY: str = "dummy"


class ChemistryMetricTest(TestCase):
    def test_ChemistryMetric(self) -> None:
        # basic test
        read_csv = pd.read_csv
        for problem_type in (
            ChemistryProblemType.DIRECT_ARYLATION,
            ChemistryProblemType.SUZUKI,
        ):
            with mock.patch(
                "ax.metrics.chemistry.pd.read_csv",
                wraps=lambda filename, index_col: read_csv(
                    filename, index_col=index_col, nrows=1
                ),
            ) as mock_read_csv:
                metric = ChemistryMetric(name="test_metric", problem_type=problem_type)
                self.assertFalse(metric.noiseless)
                self.assertIs(metric.problem_type, problem_type)
                self.assertFalse(metric.lower_is_better)
                if problem_type is ChemistryProblemType.DIRECT_ARYLATION:
                    param_names = [
                        "Base_SMILES",
                        "Concentration",
                        "Ligand_SMILES",
                        "Solvent_SMILES",
                        "Temp_C",
                    ]
                    param_values = (
                        "O=C([O-])C.[K+]",
                        0.1,
                        (
                            "CC(C)C1=CC(C(C)C)=C(C(C(C)C)=C1)C2=C(P(C3CCCCC3)"
                            "C4CCCCC4)C(OC)=CC=C2OC"
                        ),
                        "CC(N(C)C)=O",
                        105,
                    )
                    obj = 5.47
                else:
                    param_names = [
                        "Base_SMILES",
                        "Electrophile_SMILES",
                        "Ligand_SMILES",
                        "Nucleophile_SMILES",
                        "Solvent_SMILES",
                    ]
                    param_values = (
                        "[Na+].[OH-]",
                        "ClC1=CC=C(N=CC=C2)C2=C1",
                        "CC(P(C(C)(C)C)C(C)(C)C)(C)C",
                        "CC1=CC=C(N(C2CCCCO2)N=C3)C3=C1B(O)O",
                        "N#CC",
                    )
                    obj = 4.76

                params = dict(zip(param_names, param_values))
                trial = get_trial()
                trial._generator_run = GeneratorRun(
                    # pyre-fixme[6]: For 2nd argument expected `Dict[str,
                    #  Union[None, bool, float, int, str]]` but got `Dict[str,
                    #  Union[float, int, str]]`.
                    arms=[Arm(name="0_0", parameters=params)]
                )
                df = metric.fetch_trial_data(trial).unwrap().df
                self.assertEqual(mock_read_csv.call_count, 1)
                self.assertEqual(df["mean"].values[0], obj)
                self.assertTrue(np.isnan(df["sem"].values[0]))
                # test caching
                metric.fetch_trial_data(trial)
                self.assertEqual(mock_read_csv.call_count, 1)

                # test noiseless
                metric = ChemistryMetric(
                    name="test_metric", problem_type=problem_type, noiseless=True
                )
                df = metric.fetch_trial_data(trial).unwrap().df
                self.assertEqual(df["sem"].values[0], 0.0)
