import unittest

import ibm_boto3
from ibm_watson_machine_learning.preprocessing import DataJoinGraph
from ibm_watson_machine_learning.helpers.connections import DataConnection, ContainerLocation, FSLocation
from ibm_watson_machine_learning.tests.utils import is_cp4d, save_data_to_container
from ibm_watson_machine_learning.tests.autoai.abstract_tests_classes import AbstractTestOBM, \
    AbstractTestBatch

from ibm_watson_machine_learning.utils.autoai.errors import WrongDataJoinGraphNodeName

from ibm_watson_machine_learning.utils.autoai.enums import PredictionType, Metrics, RegressionAlgorithms


class TestAutoAIRemote(AbstractTestOBM, unittest.TestCase):
    """
    The test can be run on CLOUD, and CPD
    """
    HISTORICAL_RUNS_CHECK = False
    cos_resource = None

    input_data_filenames = ["go_1k.csv", "go_daily_sales.csv", "go_methods.csv",
                            "go_products.csv", "go_retailers.csv"]

    input_data_path = './autoai/data/go_sales_1k/'

    input_node_names = [name.replace('go_', 'node_').split('.')[0] for name in input_data_filenames]

    cos_endpoint = "https://s3.us.cloud-object-storage.appdomain.cloud"
    data_cos_path = "data/"

    SPACE_ONLY = False

    OPTIMIZER_NAME = "Go Sales OBM test sdk"

    target_space_id: str = None

    experiment_info = dict(
        name=OPTIMIZER_NAME,
        desc='test description',
        prediction_type=PredictionType.REGRESSION,
        prediction_column='Quantity',
        scoring=Metrics.ROOT_MEAN_SQUARED_LOG_ERROR,
        include_only_estimators=[RegressionAlgorithms.SnapDT, RegressionAlgorithms.RIDGE]
    )

    def test_01_create_multiple_data_connections__connections_created(self):
        if self.SPACE_ONLY:
            self.wml_client.set.default_space(self.space_id)
        else:
            self.wml_client.set.default_project(self.project_id)

        TestAutoAIRemote.data_connections = []

        for file, node_name in zip(self.input_data_filenames, self.input_node_names):
            asset_details = self.wml_client.data_assets.create(
                name=file,
                file_path=self.input_data_path + file)
            asset_id = asset_details['metadata']['guid']

            conn = DataConnection(
                data_join_node_name=node_name,
                data_asset_id=asset_id)

            TestAutoAIRemote.data_connections.append(conn)

        TestAutoAIRemote.results_connection = None

        self.assertIsNotNone(obj=TestAutoAIRemote.data_connections)
        self.assertGreater(len(TestAutoAIRemote.data_connections), 0)
        self.assertIsNone(self.results_connection)

    def test_02_create_data_join_graph__graph_created(self):
        data_join_graph = DataJoinGraph()
        data_join_graph.node(name="node_daily_sales")
        data_join_graph.node(name="node_methods")
        data_join_graph.node(name="node_retailers")
        data_join_graph.node(name="node_products")
        data_join_graph.node(name="node_1k", main=True)

        data_join_graph.edge(from_node="node_products", to_node="node_1k",
                             from_column=["Product number"], to_column=["Product number"])
        data_join_graph.edge(from_node="node_retailers", to_node="node_1k",
                             from_column=["Retailer code"], to_column=["Retailer code"])
        data_join_graph.edge(from_node="node_methods", to_node="node_daily_sales",
                             from_column=["Order method code"], to_column=["Order method code"])

        TestAutoAIRemote.data_join_graph = data_join_graph

        print(f"data_join_graph: {data_join_graph}")




if __name__ == '__main__':
    unittest.main()
