#!/usr/bin/env python
"""Example usage and tests for :mod:`orion.core.io.experiment_builder`."""
from __future__ import annotations

import copy
import datetime
import logging
from pathlib import Path

import pytest

import orion.core
from orion.algo.base import BaseAlgorithm
from orion.algo.random import Random
from orion.algo.space import Space
from orion.algo.tpe import TPE
from orion.core.evc.adapters import BaseAdapter
from orion.core.io import experiment_builder
from orion.core.io.config import ConfigurationError
from orion.core.io.database.ephemeraldb import EphemeralDB
from orion.core.io.database.pickleddb import PickledDB
from orion.core.io.space_builder import SpaceBuilder
from orion.core.utils import backward
from orion.core.utils.exceptions import (
    BranchingEvent,
    NoConfigurationError,
    RaceCondition,
    UnsupportedOperation,
)
from orion.core.worker.algo_wrappers import AlgoWrapper
from orion.core.worker.warm_start import KnowledgeBase
from orion.storage.base import setup_storage
from orion.storage.legacy import Legacy
from orion.testing import OrionState

def count_experiments():
    """Count experiments in storage"""
    return len(setup_storage().fetch_experiments({}))

@pytest.fixture
def space():
    """Build a space definition"""
    return {"x": "uniform(-50,50)"}

@pytest.fixture()
def python_api_config():
    """Create a configuration without the cli fluff."""
    new_config = dict(
        name="supernaekei",
        version=1,
        space={"x": "uniform(0,10)"},
        metadata={
            "user": "tsirif",
            "orion_version": "XYZ",
            "VCS": {
                "type": "git",
                "is_dirty": False,
                "HEAD_sha": "test",
                "active_branch": None,
                "diff_sha": "diff",
            },
        },
        max_trials=1000,
        max_broken=5,
        working_dir="",
        algorithm={
            "dumbalgo": {
                "done": False,
                "judgement": None,
                "scoring": 0,
                "seed": None,
                "suspend": False,
                "value": 5,
            }
        },
        _id="fasdfasfa",
        something_to_be_ignored="asdfa",
        refers=dict(root_id="supernaekei", parent_id=None, adapter=[]),
    )

    return new_config

@pytest.fixture()
def algo_unavailable_config(python_api_config):
    python_api_config["algorithm"] = {"idontreallyexist": {"but": "iwishiwould"}}
    return python_api_config

@pytest.fixture()
def new_config(random_dt, script_path):
    """Create a configuration that will not hit the database."""
    new_config = dict(
        name="supernaekei",
        metadata={
            "user": "tsirif",
            "orion_version": "XYZ",
            "user_script": script_path,
            "user_script_config": "config",
            "user_config": "abs_path/hereitis.yaml",
            "user_args": [script_path, "--mini-batch~uniform(32, 256, discrete=True)"],
            "VCS": {
                "type": "git",
                "is_dirty": False,
                "HEAD_sha": "test",
                "active_branch": None,
                "diff_sha": "diff",
            },
        },
        version=1,
        max_trials=1000,
        max_broken=5,
        working_dir="",
        algorithm={
            "dumbalgo": {
                "done": False,
                "judgement": None,
                "scoring": 0,
                "seed": None,
                "suspend": False,
                "value": 5,
            }
        },
        # attrs starting with '_' also
        _id="fasdfasfa",
        # and in general anything which is not in Experiment's slots
        something_to_be_ignored="asdfa",
        refers=dict(root_id="supernaekei", parent_id=None, adapter=[]),
    )

    backward.populate_space(new_config)

    return new_config

@pytest.fixture
def parent_version_config(script_path):
    """Return a configuration for an experiment."""
    config = dict(
        _id="parent_config",
        name="old_experiment",
        version=1,
        algorithm="random",
        metadata={
            "user": "corneauf",
            "datetime": datetime.datetime.utcnow(),
            "user_args": ["--x~normal(0,1)"],
            "user_script": script_path,
            "VCS": {
                "type": "git",
                "is_dirty": False,
                "HEAD_sha": "test",
                "active_branch": None,
                "diff_sha": "diff",
            },
            "orion_version": "XYZ",
        },
    )

    backward.populate_space(config)

    return config

@pytest.fixture
def child_version_config(parent_version_config):
    """Return a configuration for an experiment."""
    config = copy.deepcopy(parent_version_config)
    config["_id"] = "child_config"
    config["version"] = 2
    config["refers"] = {"parent_id": "parent_config"}
    config["metadata"]["datetime"] = datetime.datetime.utcnow()
    config["metadata"]["user_args"].append("--y~+normal(0,1)")
    backward.populate_space(config)
    return config

@pytest.mark.usefixtures("with_user_tsirif", "version_XYZ")
def test_get_cmd_config(raw_config):
    """Test local config (cmdconfig, cmdargs)"""
    cmdargs = {"config": raw_config}
    local_config = experiment_builder.get_cmd_config(cmdargs)

    assert local_config["algorithm"] == "random"
    assert local_config["max_trials"] == 100
    assert local_config["max_broken"] == 5
    assert local_config["name"] == "voila_voici"
    assert local_config["storage"] == {
        "database": {
            "host": "${FILE}",
            "name": "orion_test",
            "type": "pickleddb",
        }
    }
    assert local_config["metadata"] == {"orion_version": "XYZ", "user": "tsirif"}

@pytest.mark.usefixtures("with_user_tsirif", "version_XYZ")
def test_get_cmd_config_from_incomplete_config(incomplete_config_file):
    """Test local config with incomplete user configuration file
    (default, env_vars, cmdconfig, cmdargs)

    This is to ensure merge_configs update properly the subconfigs
    """
    cmdargs = {"config": incomplete_config_file}
    local_config = experiment_builder.get_cmd_config(cmdargs)

    assert "algorithm" not in local_config
    assert "max_trials" not in local_config
    assert "max_broken" not in local_config
    assert "name" not in local_config["storage"]["database"]
    assert (
        local_config["storage"]["database"]["host"] == "mongodb://user:pass@localhost"
    )
    assert local_config["storage"]["database"]["type"] == "incomplete"
    assert local_config["name"] == "incomplete"
    assert local_config["metadata"] == {"orion_version": "XYZ", "user": "tsirif"}

def test_fetch_config_from_db_no_hit():
    """Verify that fetch_config_from_db returns an empty dict when the experiment is not in db"""
    with OrionState(experiments=[], trials=[]) as cfg:
        db_config = experiment_builder.ExperimentBuilder(
            storage=cfg.storage_config
        ).fetch_config_from_db(name="supernaekei")

    assert db_config == {}

@pytest.mark.usefixtures("with_user_tsirif")
def test_fetch_config_from_db_hit(new_config):
    """Verify db config when experiment is in db"""
    with OrionState(experiments=[new_config], trials=[]) as cfg:
        db_config = experiment_builder.ExperimentBuilder(
            storage=cfg.storage_config
        ).fetch_config_from_db(name="supernaekei")

    assert db_config["name"] == new_config["name"]
    assert db_config["refers"] == new_config["refers"]
    assert db_config["metadata"] == new_config["metadata"]
    assert db_config["max_trials"] == new_config["max_trials"]
    assert db_config["max_broken"] == new_config["max_broken"]
    assert db_config["algorithm"] == new_config["algorithm"]
    assert db_config["metadata"] == new_config["metadata"]

@pytest.mark.usefixtures("with_user_tsirif")
def test_get_from_args_no_hit(config_file):
    """Try building experiment view when not in db"""
    cmdargs = {"name": "supernaekei", "config": config_file}

    with OrionState(experiments=[], trials=[]):
        with pytest.raises(NoConfigurationError) as exc_info:
            experiment_builder.get_from_args(cmdargs)
        assert "No experiment with given name 'supernaekei' and version '*'" in str(
            exc_info.value
        )

@pytest.mark.usefixtures("with_user_tsirif")
def test_get_from_args_hit(monkeypatch, raw_config, random_dt, new_config):
    """Try building experiment view when in db"""
    cmdargs = {"name": "supernaekei", "config": raw_config}

    with OrionState(experiments=[new_config], trials=[]) as cfg:
        # This is necessary because storage is instantiated inside
        # `get_from_args` with its own config not the global config set by OrionState
        def get_storage(*args, **kwargs):
            return cfg.storage

        monkeypatch.setattr(experiment_builder, "setup_storage", get_storage)

        exp_view = experiment_builder.get_from_args(cmdargs)

    assert exp_view._id == new_config["_id"]
    assert exp_view.name == new_config["name"]
    assert exp_view.configuration["refers"] == new_config["refers"]
    assert exp_view.metadata == new_config["metadata"]
    assert exp_view.max_trials == new_config["max_trials"]
    assert exp_view.max_broken == new_config["max_broken"]
    assert exp_view.algorithm
    assert exp_view.algorithm.configuration == new_config["algorithm"]

@pytest.mark.usefixtures("with_user_tsirif")
def test_get_from_args_hit_no_conf_file(
    monkeypatch, config_file, random_dt, new_config
):
    """Try building experiment view when in db, and local config file of user script does
    not exist
    """
    cmdargs = {"name": "supernaekei", "config": config_file}
    new_config["metadata"]["user_args"] += [
        "--config",
        new_config["metadata"]["user_config"],
    ]

    with OrionState(experiments=[new_config], trials=[]) as cfg:

        # This is necessary because storage is instantiated inside
        # `get_from_args` with its own config not the global config set by OrionState
        def get_storage(*args, **kwargs):
            return cfg.storage

        monkeypatch.setattr(experiment_builder, "setup_storage", get_storage)

        exp_view = experiment_builder.get_from_args(cmdargs)

    assert exp_view._id == new_config["_id"]
    assert exp_view.name == new_config["name"]
    assert exp_view.configuration["refers"] == new_config["refers"]
    assert exp_view.metadata == new_config["metadata"]
    assert exp_view.max_trials == new_config["max_trials"]
    assert exp_view.max_broken == new_config["max_broken"]
    assert exp_view.algorithm
    assert exp_view.algorithm.configuration == new_config["algorithm"]

@pytest.mark.usefixtures("with_user_dendi")
def test_build_from_args_no_hit(
    monkeypatch, config_file, random_dt, script_path, new_config
):
    """Try building experiment when not in db"""
    cmdargs = {
        "name": "supernaekei",
        "config": config_file,
        "user_args": [script_path, "x~uniform(0,10)"],
    }

    with OrionState(experiments=[], trials=[]) as cfg:
        with pytest.raises(NoConfigurationError) as exc_info:
            experiment_builder.get_from_args(cmdargs)
        assert "No experiment with given name 'supernaekei' and version '*'" in str(
            exc_info.value
        )

        # This is necessary because storage is instantiated inside
        # `get_from_args` with its own config not the global config set by OrionState
        def get_storage(*args, **kwargs):
            return cfg.storage

        monkeypatch.setattr(experiment_builder, "setup_storage", get_storage)

        exp = experiment_builder.build_from_args(cmdargs)

        assert exp.name == cmdargs["name"]
        assert exp.configuration["refers"] == {
            "adapter": [],
            "parent_id": None,
            "root_id": exp._id,
        }
        assert exp.metadata["datetime"] == random_dt
        assert exp.metadata["user"] == "dendi"
        assert exp.metadata["user_script"] == cmdargs["user_args"][0]
        assert exp.metadata["user_args"] == cmdargs["user_args"]
        assert exp.max_trials == 100
        assert exp.max_broken == 5
        assert exp.algorithm
        assert exp.algorithm.configuration == {"random": {"seed": None}}

@pytest.mark.usefixtures(
    "version_XYZ", "with_user_tsirif", "mock_infer_versioning_metadata"
)
def test_build_from_args_hit(monkeypatch, old_config_file, script_path, new_config):
    """Try building experiment when in db (no branch)"""
    cmdargs = {
        "name": "supernaekei",
        "config": old_config_file,
        "user_args": [script_path, "--mini-batch~uniform(32, 256, discrete=True)"],
    }

    with OrionState(experiments=[new_config], trials=[]) as cfg:

        # This is necessary because storage is instantiated inside
        # `get_from_args` with its own config not the global config set by OrionState
        def get_storage(*args, **kwargs):
            return cfg.storage

        monkeypatch.setattr(experiment_builder, "setup_storage", get_storage)

        # Test that experiment already exists
        experiment_builder.get_from_args(cmdargs)

        exp = experiment_builder.build_from_args(cmdargs)

    assert exp._id == new_config["_id"]
    assert exp.name == new_config["name"]
    assert exp.version == 1
    assert exp.configuration["refers"] == new_config["refers"]
    assert exp.metadata == new_config["metadata"]
    assert exp.max_trials == new_config["max_trials"]
    assert exp.max_broken == new_config["max_broken"]
    assert exp.algorithm
    assert exp.algorithm.configuration == new_config["algorithm"]

@pytest.mark.usefixtures("with_user_bouthilx")
def test_build_from_args_force_user(new_config):
    """Try building experiment view when in db"""
    cmdargs = {"name": new_config["name"]}
    cmdargs["user"] = "tsirif"
    with OrionState(experiments=[new_config], trials=[]):
        # Test that experiment already exists
        exp_view = experiment_builder.build_from_args(cmdargs)
    assert exp_view.metadata["user"] == "tsirif"

def test_build_from_args_debug_mode(monkeypatch, script_path, storage):
    """Try building experiment in debug mode"""

    experiment = experiment_builder.build_from_args(
        {
            "name": "whatever",
            "user_args": [script_path, "--mini-batch~uniform(32, 256)"],
        }
    )

    storage = experiment._storage
    assert isinstance(storage, Legacy)
    assert isinstance(storage._db, PickledDB)

    experiment = experiment_builder.build_from_args(
        {
            "name": "whatever",
            "user_args": [script_path, "--mini-batch~uniform(32, 256)"],
            "debug": True,
        }
    )

    storage = experiment._storage
    assert isinstance(storage, Legacy)
    assert isinstance(storage._db, EphemeralDB)

storage_instance = ""

def test_get_from_args_debug_mode(monkeypatch, script_path, storage):
    """Try building experiment view in debug mode"""

    old_factory = experiment_builder.setup_storage

    def retrieve_storage(*args, **kwargs):
        global storage_instance

        storage_instance = old_factory(*args, **kwargs)
        return storage_instance

    monkeypatch.setattr(experiment_builder, "setup_storage", retrieve_storage)

    # Can't build view if none exist. It's fine we only want to test the storage creation.
    with pytest.raises(NoConfigurationError):
        experiment_builder.get_from_args({"name": "whatever"})

    assert isinstance(storage_instance._db, PickledDB)

    # Can't build view if none exist. It's fine we only want to test the storage creation.
    with pytest.raises(NoConfigurationError):
        experiment_builder.get_from_args({"name": "whatever", "debug": True})

    assert isinstance(storage_instance._db, EphemeralDB)

@pytest.mark.usefixtures("with_user_tsirif", "version_XYZ")
def test_build_no_hit(config_file, random_dt, script_path):
    """Try building experiment from config when not in db"""
    name = "supernaekei"
    space = {"x": "uniform(0, 10)"}
    max_trials = 100
    max_broken = 5

    with OrionState(experiments=[], trials=[]):

        with pytest.raises(NoConfigurationError) as exc_info:
            experiment_builder.load(name)
        assert "No experiment with given name 'supernaekei' and version '*'" in str(
            exc_info.value
        )

        exp = experiment_builder.build(
            name, space=space, max_trials=max_trials, max_broken=max_broken
        )

        assert exp.name == name
        assert exp.configuration["refers"] == {
            "adapter": [],
            "parent_id": None,
            "root_id": exp._id,
        }
        assert exp.metadata == {
            "datetime": random_dt,
            "user": "tsirif",
            "orion_version": "XYZ",
        }
        assert exp.configuration["space"] == space
        assert exp.max_trials == max_trials
        assert exp.max_broken == max_broken
        assert not exp.is_done
        assert exp.algorithm
        assert exp.algorithm.configuration == {"random": {"seed": None}}

def test_build_no_commandline_config():
    """Try building experiment with no commandline configuration."""
    with OrionState(experiments=[], trials=[]):
        with pytest.raises(NoConfigurationError):
            experiment_builder.build("supernaekei")

@pytest.mark.usefixtures(
    "with_user_tsirif", "mock_infer_versioning_metadata", "version_XYZ"
)
def test_build_hit(python_api_config):
    """Try building experiment from config when in db (no branch)"""
    name = "supernaekei"

    with OrionState(experiments=[python_api_config], trials=[]):

        # Test that experiment already exists (this should fail otherwise)
        experiment_builder.load(name=name)

        exp = experiment_builder.build(**python_api_config)

    assert exp._id == python_api_config["_id"]
    assert exp.name == python_api_config["name"]
    assert exp.configuration["refers"] == python_api_config["refers"]
    assert exp.metadata == python_api_config["metadata"]
    assert exp.max_trials == python_api_config["max_trials"]
    assert exp.max_broken == python_api_config["max_broken"]
    assert exp.algorithm
    assert exp.algorithm.configuration == python_api_config["algorithm"]

@pytest.mark.usefixtures("with_user_tsirif", "version_XYZ")
def test_build_without_config_hit(python_api_config):
    """Try building experiment without commandline config when in db (no branch)"""
    name = "supernaekei"

    with OrionState(experiments=[python_api_config], trials=[]):

        # Test that experiment already exists (this should fail otherwise)
        experiment_builder.load(name=name)

        exp = experiment_builder.build(name=name)

    assert exp._id == python_api_config["_id"]
    assert exp.name == python_api_config["name"]
    assert exp.configuration["refers"] == python_api_config["refers"]
    assert exp.metadata == python_api_config["metadata"]
    assert exp.max_trials == python_api_config["max_trials"]
    assert exp.max_broken == python_api_config["max_broken"]
    assert exp.algorithm
    assert exp.algorithm.configuration == python_api_config["algorithm"]

@pytest.mark.usefixtures(
    "with_user_tsirif", "version_XYZ", "mock_infer_versioning_metadata"
)
def test_build_from_args_without_cmd(
    monkeypatch, old_config_file, script_path, new_config
):
    """Try building experiment without commandline when in db (no branch)"""
    name = "supernaekei"

    cmdargs = {"name": name, "config": old_config_file}

    with OrionState(experiments=[new_config], trials=[]) as cfg:

        # This is necessary because storage is instantiated inside
        # `get_from_args` with its own config not the global config set by OrionState
        def get_storage(*args, **kwargs):
            return cfg.storage

        monkeypatch.setattr(experiment_builder, "setup_storage", get_storage)

        # Test that experiment already exists (this should fail otherwise)
        experiment_builder.get_from_args(cmdargs)

        exp = experiment_builder.build_from_args(cmdargs)

    assert exp._id == new_config["_id"]
    assert exp.name == new_config["name"]
    assert exp.configuration["refers"] == new_config["refers"]
    assert exp.metadata == new_config["metadata"]
    assert exp.max_trials == new_config["max_trials"]
    assert exp.max_broken == new_config["max_broken"]
    assert exp.algorithm
    assert exp.algorithm.configuration == new_config["algorithm"]

# TODO: Remove for v0.4
class TestStrategyDeprecated:
    def test_strategy_not_defined(self, caplog, space):
        """Verify there is no warning"""
        with OrionState():
            with caplog.at_level(logging.WARNING):
                exp = experiment_builder.build(name="whatever", space=space)
            assert "`strategy` option is not supported anymore." not in caplog.text

    def test_strategy_defined_in_global_config(self, caplog, space, monkeypatch):
        """Verify there is a warning"""

        with monkeypatch.context() as m:
            m.setattr(
                orion.core.config.experiment,
                "strategy",
                {"this is deprecated": "and should be ignored"},
            )
            with OrionState():
                with caplog.at_level(logging.WARNING):
                    exp = experiment_builder.build(name="whatever", space=space)
                assert "`strategy` option is not supported anymore." in caplog.text

    def test_strategy_defined_in_config(self, caplog, space):
        """Verify there is a warning"""
        with OrionState():
            with caplog.at_level(logging.WARNING):
                exp = experiment_builder.build(
                    name="whatever",
                    space=space,
                    strategy={"this is deprecated": "and should be ignored"},
                )
            assert "`strategy` option is not supported anymore." in caplog.text

@pytest.mark.usefixtures(
    "with_user_tsirif", "version_XYZ", "mock_infer_versioning_metadata"
)
class TestExperimentVersioning:
    """Create new Experiment with auto-versioning."""

    def test_new_experiment_wout_version(self, space):
        """Create a new and never-seen-before experiment without a version."""
        with OrionState():
            exp = experiment_builder.build(name="exp_wout_version", space=space)

        assert exp.version == 1

    def test_new_experiment_w_version(self, space):
        """Create a new and never-seen-before experiment with a version."""
        with OrionState():
            exp = experiment_builder.build(
                name="exp_wout_version", version=1, space=space
            )

        assert exp.version == 1

    def test_experiment_overwritten_evc_disabled(self, parent_version_config, caplog):
        """Build an existing experiment with different config, overwriting previous config."""
        parent_version_config.pop("version")

        # Note: this is a bit of a hack, but we need to make sure the config is also compatible with gridsearch.
        parent_version_config = copy.deepcopy(parent_version_config)
        parent_version_config["space"] = {"/x": "uniform(0, 1)"}
        with OrionState(experiments=[parent_version_config]):

            with caplog.at_level(logging.WARNING):

                exp = experiment_builder.build(name=parent_version_config["name"])
                assert "Running experiment in a different state" not in caplog.text

            assert exp.version == 1
            assert exp.configuration["algorithm"] == {"random": {"seed": None}}

            new_algo = "gridsearch"
            with caplog.at_level(logging.WARNING):

                exp = experiment_builder.build(
                    name=parent_version_config["name"], algorithm=new_algo
                )
                assert "Running experiment in a different state" in caplog.text

            assert exp.version == 1
            assert list(exp.configuration["algorithm"].keys())[0] == new_algo

            caplog.clear()
            with caplog.at_level(logging.WARNING):

                exp = experiment_builder.load(name=parent_version_config["name"])
                assert "Running experiment in a different state" not in caplog.text

            assert exp.version == 1
            assert list(exp.configuration["algorithm"].keys())[0] == new_algo

    def test_backward_compatibility_no_version(self, parent_version_config):
        """Branch from parent that has no version field."""
        parent_version_config.pop("version")
        with OrionState(experiments=[parent_version_config]):
            exp = experiment_builder.build(
                name=parent_version_config["name"],
                space={"y": "uniform(0, 10)"},
                branching={"enable": True},
            )

        assert exp.version == 2

    @pytest.mark.usefixtures("mock_infer_versioning_metadata")
    def test_old_experiment_wout_version(self, parent_version_config):
        """Create an already existing experiment without a version."""
        with OrionState(experiments=[parent_version_config]):
            exp = experiment_builder.build(name=parent_version_config["name"])

        assert exp.version == 1

    @pytest.mark.usefixtures("mock_infer_versioning_metadata")
    def test_old_experiment_2_wout_version(
        self, parent_version_config, child_version_config
    ):
        """Create an already existing experiment without a version and getting last one."""
        with OrionState(experiments=[parent_version_config, child_version_config]):
            exp = experiment_builder.build(name=parent_version_config["name"])

        assert exp.version == 2

    @pytest.mark.usefixtures("mock_infer_versioning_metadata")
    def test_old_experiment_w_version(
        self, parent_version_config, child_version_config
    ):
        """Create an already existing experiment with a version."""
        with OrionState(experiments=[parent_version_config, child_version_config]):
            exp = experiment_builder.build(
                name=parent_version_config["name"], version=1
            )

        assert exp.version == 1

    @pytest.mark.usefixtures("mock_infer_versioning_metadata")
    def test_old_experiment_w_version_bigger_than_max(
        self, parent_version_config, child_version_config
    ):
        """Create an already existing experiment with a too large version."""
        with OrionState(experiments=[parent_version_config, child_version_config]):
            exp = experiment_builder.build(
                name=parent_version_config["name"], version=8
            )

        assert exp.version == 2

@pytest.mark.usefixtures("with_user_tsirif", "version_XYZ")
class TestBuild:
    """Test building the experiment"""

    @pytest.mark.usefixtures("mock_infer_versioning_metadata")
    def test_good_set_before_init_hit_no_diffs_exc_max_trials(self, new_config):
        """Trying to set, and NO differences were found from the config pulled from db.

        Everything is normal, nothing changes. Experiment is resumed,
        perhaps with more trials to evaluate (an exception is 'max_trials').
        """
        with OrionState(experiments=[new_config], trials=[]):

            new_config["max_trials"] = 5000

            exp = experiment_builder.build(**new_config)

        # Deliver an external configuration to finalize init
        new_config["algorithm"]["dumbalgo"]["done"] = False
        new_config["algorithm"]["dumbalgo"]["judgement"] = None
        new_config["algorithm"]["dumbalgo"]["scoring"] = 0
        new_config["algorithm"]["dumbalgo"]["suspend"] = False
        new_config["algorithm"]["dumbalgo"]["value"] = 5
        new_config["algorithm"]["dumbalgo"]["seed"] = None
        new_config.pop("something_to_be_ignored")
        new_config["knowledge_base"] = None
        assert exp.configuration == new_config

    @pytest.mark.usefixtures("mock_infer_versioning_metadata")
    def test_good_set_before_init_no_hit(self, random_dt, new_config):
        """Trying to set, overwrite everything from input."""
        with OrionState(experiments=[], trials=[]):
            exp = experiment_builder.build(**new_config)
            found_config = list(
                setup_storage().fetch_experiments(
                    {"name": "supernaekei", "metadata.user": "tsirif"}
                )
            )

        new_config["metadata"]["datetime"] = exp.metadata["datetime"]

        assert len(found_config) == 1
        _id = found_config[0].pop("_id")
        assert _id != "fasdfasfa"
        assert exp._id == _id
        new_config["refers"] = {}
        new_config.pop("_id")
        new_config.pop("something_to_be_ignored")
        new_config["algorithm"]["dumbalgo"]["done"] = False
        new_config["algorithm"]["dumbalgo"]["judgement"] = None
        new_config["algorithm"]["dumbalgo"]["scoring"] = 0
        new_config["algorithm"]["dumbalgo"]["suspend"] = False
        new_config["algorithm"]["dumbalgo"]["value"] = 5
        new_config["algorithm"]["dumbalgo"]["seed"] = None
        new_config["refers"] = {"adapter": [], "parent_id": None, "root_id": _id}
        new_config["knowledge_base"] = None
        assert found_config[0] == new_config
        assert exp.name == new_config["name"]
        assert exp.configuration["refers"] == new_config["refers"]
        assert exp.metadata == new_config["metadata"]
        assert exp.max_trials == new_config["max_trials"]
        assert exp.max_broken == new_config["max_broken"]
        assert exp.working_dir == new_config["working_dir"]
        assert exp.version == new_config["version"]
        assert exp.algorithm
        assert exp.algorithm.configuration == new_config["algorithm"]

    def test_working_dir_is_correctly_set(self, new_config):
        """Check if working_dir is correctly changed."""
        with OrionState():
            new_config["working_dir"] = "./"
            exp = experiment_builder.build(**new_config)
            storage = setup_storage()
            found_config = list(
                storage.fetch_experiments(
                    {"name": "supernaekei", "metadata.user": "tsirif"}
                )
            )

            found_config = found_config[0]
            exp = experiment_builder.build(**found_config)
            assert exp.working_dir == "./"

    def test_working_dir_works_when_db_absent(self, database, new_config):
        """Check if working_dir is correctly when absent from the database."""
        with OrionState(experiments=[], trials=[]):
            exp = experiment_builder.build(**new_config)
            storage = setup_storage()
            found_config = list(
                storage.fetch_experiments(
                    {"name": "supernaekei", "metadata.user": "tsirif"}
                )
            )

            found_config = found_config[0]
            exp = experiment_builder.build(**found_config)
            assert exp.working_dir == ""

    @pytest.mark.usefixtures("mock_infer_versioning_metadata")
    def test_configuration_hit_no_diffs(self, new_config):
        """Return a configuration dict according to an experiment object.

        Before initialization is done, it can be the case that the pair (`name`,
        user's name) has not hit the database. return a yaml compliant form
        of current state, to be used with :mod:`orion.core.cli.esolve_config`.
        """
        with OrionState(experiments=[new_config], trials=[]):
            experiment_count_before = count_experiments()
            exp = experiment_builder.build(**new_config)
            assert experiment_count_before == count_experiments()

        new_config["algorithm"]["dumbalgo"]["done"] = False
        new_config["algorithm"]["dumbalgo"]["judgement"] = None
        new_config["algorithm"]["dumbalgo"]["scoring"] = 0
        new_config["algorithm"]["dumbalgo"]["suspend"] = False
        new_config["algorithm"]["dumbalgo"]["value"] = 5
        new_config["algorithm"]["dumbalgo"]["seed"] = None
        new_config.pop("something_to_be_ignored")
        new_config["knowledge_base"] = None
        assert exp.configuration == new_config

    def test_instantiation_after_init(self, new_config):
        """Verify that algo, space and refers was instantiated properly"""
        with OrionState(experiments=[new_config], trials=[]):
            exp = experiment_builder.build(**new_config)

        assert isinstance(exp.algorithm, AlgoWrapper)
        assert isinstance(exp.space, Space)
        assert isinstance(exp.refers["adapter"], BaseAdapter)

    @pytest.mark.usefixtures("mock_infer_versioning_metadata")
    def test_algo_case_insensitive(self, new_config):
        """Verify that algo with uppercase or lowercase leads to same experiment"""
        with OrionState(experiments=[new_config], trials=[]):
            new_config["algorithm"]["DUMBALGO"] = new_config["algorithm"].pop(
                "dumbalgo"
            )
            exp = experiment_builder.build(**new_config)

            assert exp.version == 1

    def test_hierarchical_space(self, new_config):
        """Verify space can have hierarchical structure"""
        space = {
            "a": {"x": "uniform(0, 10, discrete=True)"},
            "b": {"y": "loguniform(1e-08, 1)", "z": "choices(['voici', 'voila', 2])"},
        }

        with OrionState(experiments=[], trials=[]):
            exp = experiment_builder.build("hierarchy", space=space)

            exp2 = experiment_builder.build("hierarchy")

        assert "a.x" in exp.space
        assert "b.y" in exp.space
        assert "b.z" in exp.space

        # Make sure it can be fetched properly from db as well
        assert "a.x" in exp2.space
        assert "b.y" in exp2.space
        assert "b.z" in exp2.space

    def test_try_set_after_race_condition(self, new_config, monkeypatch):
        """Cannot set a configuration after init if it looses a race
        condition.

        The experiment from process which first writes to db is initialized
        properly. The experiment which looses the race condition cannot be
        initialized and needs to be rebuilt.
        """
        with OrionState(experiments=[new_config], trials=[]) as cfg:
            experiment_count_before = count_experiments()

            def insert_race_condition(*args, **kwargs):
                if insert_race_condition.count == 0:
                    data = {}
                else:
                    data = new_config

                insert_race_condition.count += 1

                return data

            insert_race_condition.count = 0

            builder = experiment_builder.ExperimentBuilder(cfg.storage)
            monkeypatch.setattr(builder, "fetch_config_from_db", insert_race_condition)

            builder.build(**new_config)

            assert experiment_count_before == count_experiments()

        # Should be called
        # - once in build(),
        #     -> then register fails,
        # - then called once again in build,
        # - then called in load to evaluate the conflicts
        assert insert_race_condition.count == 3

    def test_algorithm_config_with_just_a_string(self):
        """Test that configuring an algorithm with just a string is OK."""
        name = "supernaedo3"
        space = {"x": "uniform(0,10)"}
        algorithm = "dumbalgo"

        with OrionState(experiments=[], trials=[]):
            exp = experiment_builder.build(name=name, space=space, algorithm=algorithm)

        assert exp.configuration["algorithm"] == {
            "dumbalgo": {
                "done": False,
                "judgement": None,
                "scoring": 0,
                "suspend": False,
                "value": (5,),
                "seed": None,
            }
        }

    def test_new_child_with_branch(self):
        """Check that experiment is not incremented when branching with a new name."""
        name = "parent"
        space = {"x": "uniform(0, 10)"}

        with OrionState(experiments=[], trials=[]):
            parent = experiment_builder.build(name=name, space=space)

            assert parent.name == name
            assert parent.version == 1

            child_name = "child"

            child = experiment_builder.build(
                name=name, branching={"branch_to": child_name, "enable": True}
            )

            assert child.name == child_name
            assert child.version == 1
            assert child.refers["parent_id"] == parent.id

            child_name = "child2"

            child = experiment_builder.build(
                name=child_name, branching={"branch_from": name, "enable": True}
            )

            assert child.name == child_name
            assert child.version == 1
            assert child.refers["parent_id"] == parent.id

    def test_no_increment_when_child_exist(self):
        """Check that experiment cannot be incremented when asked for v1 while v2 exists."""
        name = "parent"
        space = {"x": "uniform(0,10)"}

        with OrionState(experiments=[], trials=[]):
            parent = experiment_builder.build(name=name, space=space)
            child = experiment_builder.build(
                name=name, space={"x": "loguniform(1,10)"}, branching={"enable": True}
            )
            assert child.name == parent.name
            assert parent.version == 1
            assert child.version == 2

            with pytest.raises(BranchingEvent) as exc_info:
                experiment_builder.build(
                    name=name,
                    version=1,
                    space={"x": "loguniform(1,10)"},
                    branching={"enable": True},
                )
            assert "Configuration is different and generates a branching" in str(
                exc_info.value
            )

    def test_race_condition_wout_version(self, monkeypatch):
        """Test that an experiment losing the race condition during version increment raises
        RaceCondition if version number was not specified.
        """
        name = "parent"
        space = {"x": "uniform(0,10)"}

        with OrionState(experiments=[], trials=[]):
            parent = experiment_builder.build(name, space=space)
            child = experiment_builder.build(
                name=name, space={"x": "loguniform(1,10)"}, branching={"enable": True}
            )
            assert child.name == parent.name
            assert parent.version == 1
            assert child.version == 2

            # Either
            # 1.
            #     fetch_config_from_db only fetch parent
            #     test_version finds other child
            #     -> Detect race condition looking at conflicts
            # 2.
            #     fetch_config_from_db only fetch parent
            #     test_version do not find other child
            #     -> DuplicateKeyError

            def insert_race_condition_1(self, query):
                is_auto_version_query = query == {
                    "name": name,
                    "refers.parent_id": parent.id,
                }
                if is_auto_version_query:
                    data = [child.configuration]
                # First time the query returns no other child
                elif insert_race_condition_1.count < 1:
                    data = [parent.configuration]
                else:
                    data = [parent.configuration, child.configuration]

                insert_race_condition_1.count += int(is_auto_version_query)

                return data

            insert_race_condition_1.count = 0

            monkeypatch.setattr(
                setup_storage().__class__, "fetch_experiments", insert_race_condition_1
            )

            with pytest.raises(RaceCondition) as exc_info:
                experiment_builder.build(
                    name=name,
                    space={"x": "loguniform(1,10)"},
                    branching={"enable": True},
                )
            assert "There was likely a race condition during version" in str(
                exc_info.value
            )

            def insert_race_condition_2(self, query):
                is_auto_version_query = query == {
                    "name": name,
                    "refers.parent_id": parent.id,
                }
                # First time the query returns no other child
                if is_auto_version_query:
                    data = []
                elif insert_race_condition_2.count < 1:
                    data = [parent.configuration]
                else:
                    data = [parent.configuration, child.configuration]

                insert_race_condition_2.count += int(is_auto_version_query)

                return data

            insert_race_condition_2.count = 0

            monkeypatch.setattr(
                setup_storage().__class__, "fetch_experiments", insert_race_condition_2
            )

            with pytest.raises(RaceCondition) as exc_info:
                experiment_builder.build(
                    name=name,
                    space={"x": "loguniform(1,10)"},
                    branching={"enable": True},
                )
            assert "There was a race condition during branching." in str(exc_info.value)

    def test_race_condition_w_version(self, monkeypatch):
        """Test that an experiment losing the race condition during version increment cannot
        be resolved automatically if a version number was specified.

        Note that if we would raise RaceCondition, the conflict would still occur since
        the version number fetched will not be the new one from the resolution but the requested
        one. Therefore raising and handling RaceCondition would lead to infinite recursion in
        the experiment builder.
        """
        name = "parent"
        space = {"x": "uniform(0,10)"}

        with OrionState(experiments=[], trials=[]):
            parent = experiment_builder.build(name, space=space)
            child = experiment_builder.build(
                name=name, space={"x": "loguniform(1,10)"}, branching={"enable": True}
            )
            assert child.name == parent.name
            assert parent.version == 1
            assert child.version == 2

            # Either
            # 1.
            #     fetch_config_from_db only fetch parent
            #     test_version finds other child
            #     -> Detect race condition looking at conflicts
            # 2.
            #     fetch_config_from_db only fetch parent
            #     test_version do not find other child
            #     -> DuplicateKeyError

            def insert_race_condition_1(self, query):
                is_auto_version_query = query == {
                    "name": name,
                    "refers.parent_id": parent.id,
                }
                if is_auto_version_query:
                    data = [child.configuration]
                # First time the query returns no other child
                elif insert_race_condition_1.count < 1:
                    data = [parent.configuration]
                else:
                    data = [parent.configuration, child.configuration]

                insert_race_condition_1.count += int(is_auto_version_query)

                return data

            insert_race_condition_1.count = 0

            monkeypatch.setattr(
                setup_storage().__class__, "fetch_experiments", insert_race_condition_1
            )

            with pytest.raises(BranchingEvent) as exc_info:
                experiment_builder.build(
                    name=name,
                    version=1,
                    space={"x": "loguniform(1,10)"},
                    branching={"enable": True},
                )
            assert "Configuration is different and generates" in str(exc_info.value)

            def insert_race_condition_2(self, query):
                is_auto_version_query = query == {
                    "name": name,
                    "refers.parent_id": parent.id,
                }
                # First time the query returns no other child
                if is_auto_version_query:
                    data = []
                elif insert_race_condition_2.count < 1:
                    data = [parent.configuration]
                else:
                    data = [parent.configuration, child.configuration]

                insert_race_condition_2.count += int(is_auto_version_query)

                return data

            insert_race_condition_2.count = 0

            monkeypatch.setattr(
                setup_storage().__class__, "fetch_experiments", insert_race_condition_2
            )

            with pytest.raises(RaceCondition) as exc_info:
                experiment_builder.build(
                    name=name,
                    version=1,
                    space={"x": "loguniform(1,10)"},
                    branching={"enable": True},
                )
            assert "There was a race condition during branching." in str(exc_info.value)

    def test_build_experiment_with_kb(self, tmp_path: Path):
        """Test that passing a configuration for the KB to `create_experiment` works."""
        exp_storage_file = str(tmp_path / "db.pkl")
        kb_storage_file = str(tmp_path / "kb.pkl")
        experiment = experiment_builder.build(
            "test",
            space={"x": "uniform(0, 10)"},
            storage={
                "type": "legacy",
                "database": {"type": "pickleddb", "host": exp_storage_file},
            },
            knowledge_base={
                KnowledgeBase.__qualname__: {
                    "storage": {
                        "type": "legacy",
                        "database": {"type": "pickleddb", "host": kb_storage_file},
                    }
                }
            },
        )
        assert experiment.knowledge_base is not None
        assert isinstance(experiment.knowledge_base, KnowledgeBase)

class TestInstantiateKB:
    def test_build_experiment_with_bad_kb_config(self):
        """Test that passing a bad configuration for the KB raises an error."""
        with pytest.raises(
            ConfigurationError,
            match="The configuration for the KB should only have one key",
        ):
            experiment_builder._instantiate_knowledge_base(
                {
                    "fooooobar": {
                        "storage": {"type": "legacy", "database": {"type": "bad"}}
                    },
                    "baz": 123,
                }
            )

    def test_kb_class_not_found(self):
        with pytest.raises(
            ConfigurationError,
            match="Unable to find a subclass of KnowledgeBase with the given name",
        ):
            experiment_builder._instantiate_knowledge_base(
                {
                    "NonExistentKB": {
                        "storage": {"type": "legacy", "database": {"type": "bad"}}
                    },
                }
            )

    def test_finds_kb_subclass_and_uses_it(self, tmp_path: Path):
        class MyKB(KnowledgeBase):
            pass

        path = tmp_path / "db.pkl"
        kb = experiment_builder._instantiate_knowledge_base(
            {
                "MyKB": {
                    "storage": {
                        "type": "legacy",
                        "database": {"type": "pickleddb", "host": str(path)},
                    }
                },
            }
        )
        assert isinstance(kb, MyKB)

    def test_multiple_subclasses_match_name(self):
        class KB(KnowledgeBase):  # noqa
            pass

        class KB(KnowledgeBase):  # noqa
            pass

        with pytest.raises(
            ConfigurationError,
            match="Multiple subclasses of KnowledgeBase with the given name",
        ):
            experiment_builder._instantiate_knowledge_base(
                {
                    "KB": {"storage": {"type": "legacy", "database": {"type": "bad"}}},
                }
            )

@pytest.fixture
def space_obj(space: dict[str, str]):
    # todo: Rename `space` above to `space_config` and rename this to just `space`.
    return SpaceBuilder().build(space)

class TestInstantiateAlgo:
    """Tests for the `_instantiate_algo` function."""

    @pytest.mark.parametrize("algo_class_name", ["Random", "TPE"])
    @pytest.mark.parametrize("lowercase", [True, False])
    @pytest.mark.parametrize("max_trials", [None, 10])
    def test_with_class_name(
        self,
        algo_class_name: str,
        space_obj: Space,
        lowercase: bool,
        max_trials: int | None,
    ):
        """Test instantiating an algorithm by passing the class name as a config."""
        algo = experiment_builder._instantiate_algo(
            space=space_obj,
            max_trials=max_trials,
            config=algo_class_name.lower() if lowercase else algo_class_name,
        )
        assert isinstance(algo, BaseAlgorithm)
        assert type(algo.unwrapped).__qualname__ == algo_class_name
        assert algo.max_trials == max_trials

    @pytest.mark.parametrize("algo_class", [Random, TPE])
    @pytest.mark.parametrize("max_trials", [None, 10])
    def test_with_algo_class(
        self, algo_class: type[BaseAlgorithm], space_obj: Space, max_trials: int | None
    ):
        """Test instantiating an algorithm by passing the class as a config."""

        algo = experiment_builder._instantiate_algo(
            space=space_obj,
            max_trials=max_trials,
            config=algo_class,
        )
        assert isinstance(algo, BaseAlgorithm)
        assert isinstance(algo.unwrapped, algo_class)
        assert algo.max_trials == max_trials

    @pytest.mark.parametrize("algo_class", [Random, TPE])
    @pytest.mark.parametrize("max_trials", [None, 10])
    def test_with_dict(
        self, algo_class: type[BaseAlgorithm], space_obj: Space, max_trials: int | None
    ):
        """Test instantiating an algorithm using a config dictionary."""
        algo = experiment_builder._instantiate_algo(
            space=space_obj,
            max_trials=max_trials,
            config={algo_class.__qualname__.lower(): {}},
        )
        assert isinstance(algo, BaseAlgorithm)
        assert isinstance(algo.unwrapped, algo_class)
        assert algo.max_trials == max_trials

def test_load_unavailable_algo(algo_unavailable_config, capsys):
    with OrionState(experiments=[algo_unavailable_config]):
        experiment = experiment_builder.load("supernaekei", mode="r")
        assert experiment.algorithm == algo_unavailable_config["algorithm"]
        assert (
            experiment.configuration["algorithm"]
            == algo_unavailable_config["algorithm"]
        )

        experiment = experiment_builder.load("supernaekei", mode="w")
        assert experiment.algorithm == algo_unavailable_config["algorithm"]
        assert (
            experiment.configuration["algorithm"]
            == algo_unavailable_config["algorithm"]
        )

        with pytest.raises(
            NotImplementedError, match="Could not find implementation of BaseAlgorithm"
        ):
            experiment_builder.build("supernaekei")

class TestInitExperimentReadWrite:
    """Create new Experiment instance that only supports read/write."""

    def test_empty_experiment_rw(self):
        """Hit user name, but exp_name does not hit the db."""
        with OrionState(experiments=[], trials=[]):
            with pytest.raises(NoConfigurationError) as exc_info:
                experiment_builder.load("supernaekei")
            assert "No experiment with given name 'supernaekei' and version '*'" in str(
                exc_info.value
            )

    def test_existing_experiment_read_mode(self, new_config):
        """Hit exp_name + user's name in the db, support read only."""
        with OrionState(experiments=[new_config], trials=[]):
            exp = experiment_builder.load(name="supernaekei", mode="r")

        assert exp.mode == "r"
        exp.fetch_trials()  # Should pass
        with pytest.raises(UnsupportedOperation) as exc:
            exp.fix_lost_trials()
        assert exc.match("to execute `fix_lost_trials()")

    def test_existing_experiment_read_write_mode(self, new_config):
        """Hit exp_name + user's name in the db, support read and write."""
        with OrionState(experiments=[new_config], trials=[]):
            exp = experiment_builder.load(name="supernaekei", mode="w")

        assert exp.mode == "w"
        exp.fetch_trials()  # Should pass
        exp.fix_lost_trials()  # Should pass

        with pytest.raises(UnsupportedOperation) as exc:
            exp.reserve_trial()
        assert exc.match("to execute `reserve_trial()")


"""
Test module for weight.py.
"""
import os
import numpy as np
import rtree
from numpy.testing import assert_array_equal, assert_almost_equal
from shapely.geometry import Polygon
from pyproj import CRS
from osgeo import ogr

from inflow import weight

TEST_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)))
DATA_DIR = os.path.join(TEST_DIR, 'data')
OUTPUT_DIR = os.path.join(TEST_DIR, 'output')
MAIN_DIR = os.path.join(TEST_DIR, os.pardir, 'inflow')

def read_weight_table(fname, ncol=7):
    """
    Read a weight-table csv file.

    Parameters
    ----------
    fname : str
        Name of file.
    ncol : int, optional
        Number of columns to read.

    Returns
    -------
    data : ndarray
        Contents of weight table as an array.
    """
    data = np.genfromtxt(fname, delimiter=',', skip_header=1)
    data = data[:,:ncol]

    return data

def compare_weight_table_files(output_file, benchmark_file, ncol=7):
    """
    Check if the contents of two weight-table files match.

    Parameters
    ----------
    output_file : str
        Name of file to compare.
    benchmark_file : str
        Name of file to compare.
    ncol : int, optional
        Number of columns to read.

    Returns
    -------
    data : ndarray
        Contents of weight table as an array.
    """
    output = read_weight_table(output_file, ncol=ncol)
    benchmark = read_weight_table(benchmark_file, ncol=ncol)
    assert_almost_equal(output, benchmark)

def test_parse_coordinate_order():
    """
    Verify that parse_coordinate_order correctly identifies coordinate order.
    """
    shapefile = os.path.join(DATA_DIR, 'catchment',
                             'thames_near_cirencester_uk',
                             'Catchment_thames_drainID45390.shp')
    catchment_file_obj = ogr.Open(shapefile)
    catchment_geospatial_layer = catchment_file_obj.GetLayer()
    catchment_spatial_reference = catchment_geospatial_layer.GetSpatialRef()
    catchment_wkt = catchment_spatial_reference.ExportToWkt()
    original_catchment_crs = CRS.from_wkt(catchment_wkt)

    catchment_coordinate_order = weight.parse_coordinate_order(
        original_catchment_crs)

    expected = 'XY'

    assert catchment_coordinate_order == expected

def test_generate_unique_id():
    """
    Verify that generate_unique_id() returns weight table ID in the correct
    format, i.e.
    <catchment shapefile ID><3-digit latitude><3-digit longitude>
    as an integer.
    """
    rivid = 12345
    lat_index = 6
    lon_index = 7

    uid = weight.generate_unique_id(rivid, lat_index, lon_index)

    expected = 1234500060007

    assert uid == expected

def test_extract_lat_lon_from_nc():
    """
    Verify that extract_lat_lon_from_nc() correctly extracts latitude and
    longitude variables from a netCDF file.
    """
    filename = os.path.join(DATA_DIR, 'lsm_grids',
                            'gldas2', 'GLDAS_NOAH025_3H.A20101231.0000.020.nc4')
    lat_variable = 'lat'
    lon_variable = 'lon'

    lat, lon = weight.extract_lat_lon_from_nc(filename,
                                              lat_variable=lat_variable,
                                              lon_variable=lon_variable)

    expected_lat_size = (600,)
    expected_lon_size =  (1440,)

    lat_lon_sizes = [lat.shape, lon.shape]
    expected = [expected_lat_size, expected_lon_size]

    assert lat_lon_sizes == expected

def test_calculate_polygon_area():
    """
    Verify that a square of 30 arc-seconds is approximately 8.5 x 10^5 m^2
    at the equator.
    """
    deg = 0.00833333 # 30 arc-seconds in degrees.
    polygon = Polygon([[0, 0], [deg, 0], [deg, deg], [0, deg]])

    area = weight.calculate_polygon_area(polygon)
    expected = 854795.9840483834

    assert area == expected

def test_shift_longitude():
    """
    Verify that sample coordinates from a 0 to 360 degree longitude
    coordinate system (e.g. the convention for some ECMWF grids) are
    correctly mapped to a -180 to 180 longitude coordinate system.
    """
    lon = np.array([180, 360])

    shifted = weight.shift_longitude(lon)
    expected = np.array([-180, 0])

    assert_array_equal(shifted, expected)

def test_reproject():
    """
    Verify that coordinates in the EPSG:4269/NAD83 coordinate reference
    system are correctly reprojected to EPSG:4326/WGS84.
    """
    lat = 38.21056641800004
    lon = -106.52128819799998

    original_wkt = ('GEOGCRS["NAD83",DATUM["North American Datum 1983",'
                    'ELLIPSOID["GRS 1980",6378137,298.257222101,'
                    'LENGTHUNIT["metre",1]]],PRIMEM["Greenwich",0,'
                    'ANGLEUNIT["degree",0.0174532925199433]],'
                    'CS[ellipsoidal,2],AXIS["geodetic latitude (Lat)",north,'
                    'ORDER[1],ANGLEUNIT["degree",0.0174532925199433]],'
                    'AXIS["geodetic longitude (Lon)",east,ORDER[2],'
                    'ANGLEUNIT["degree",0.0174532925199433]],ID["EPSG",4269]]')

    original_crs = CRS.from_wkt(original_wkt)

    geographic_crs = CRS.from_epsg(4326)

    latlon = weight.reproject(lon, lat, original_crs, geographic_crs,
                              always_xy=True)

    expected = (-106.52128597511603, 38.21056645013602)

    assert latlon == expected

def test_reproject_extent():
    """
    Verify that extents (i.e. [min lon, max lon, min lat, max lat]) in the
    EPSG:4269/NAD83 coordinate reference system are correctly reprojected
    to EPSG:4326/WGS84.
    """
    extent = [-106.52128819799998, -106.42948920899998,
              38.21056641800004, 38.32185130200008,]

    original_wkt = ('GEOGCRS["NAD83",DATUM["North American Datum 1983",'
                    'ELLIPSOID["GRS 1980",6378137,298.257222101,'
                    'LENGTHUNIT["metre",1]]],PRIMEM["Greenwich",0,'
                    'ANGLEUNIT["degree",0.0174532925199433]],'
                    'CS[ellipsoidal,2],AXIS["geodetic latitude (Lat)",north,'
                    'ORDER[1],ANGLEUNIT["degree",0.0174532925199433]],'
                    'AXIS["geodetic longitude (Lon)",east,ORDER[2],'
                    'ANGLEUNIT["degree",0.0174532925199433]],ID["EPSG",4269]]')

    original_crs = CRS.from_wkt(original_wkt)

    geographic_crs = CRS.from_epsg(4326)

    extent = weight.reproject_extent(extent, original_crs, geographic_crs,
                                     always_xy=True)

    expected = [-106.52128597511603, -106.42948675047731,
                38.21056645013602, 38.32185167233476]

    assert extent == expected

def test_generate_rtree():
    """
    Verify that an Rtree spatial index populated with a single polygon
    has bounds consistent with that polygon.
    """
    poly = Polygon([[0, 0], [1, 0], [1, 1], [0, 1], [0, 0]])

    feature_list = [{'polygon': poly}]

    tree = weight.generate_rtree(feature_list)

    bounds = tree.bounds

    expected = [0.0, 0.0, 1.0, 1.0]

    assert bounds == expected

def test_get_lat_lon_indices():
    """
    Verify that get_lat_lon_indices returns the correct indices in
    latitude/longitude arrays given sample coordinates.
    """

    lat_array = np.array([0, 45, 90])
    lon_array = np.array([0, 90, 180])

    lat = 45
    lon = 180

    latlon = weight.get_lat_lon_indices(lat_array, lon_array, lat, lon)
    expected = (1, 2)

    assert latlon == expected

def test_write_weight_table():
    """
    Verify that write_weight_table produces a weight table consistent with
    benchmark.
    """
    lsm_file = os.path.join(DATA_DIR, 'lsm_grids', 'gldas2',
                            'GLDAS_NOAH025_3H.A20101231.0000.020.nc4')
    lsm_grid_lat, lsm_grid_lon = weight.extract_lat_lon_from_nc(lsm_file)

    catchment_shapefile = os.path.join(DATA_DIR, 'catchment',
                                       'saguache_colorado', 'catchment.shp')
    shp = ogr.Open(catchment_shapefile)
    catchment_layer = shp.GetLayer()

    connect_rivid_list = [17880832]

    catchment_rivid_list = [17880282, 17880832, 17880836, 17880268,
                            17880834, 17880284, 17880830, 17880298]

    lsm_grid_polygon = Polygon([[-106.25, 38.25], [-106.5, 38.25],
                                [-106.5, 38.5], [-106.25, 38.5],
                                [-106.25, 38.25]])

    lsm_grid_voronoi_feature_list = [{
        'polygon': lsm_grid_polygon, 'lon': -106.375, 'lat': 38.375}]

    lsm_grid_rtree = rtree.index.Index()
    lsm_grid_rtree.insert(0, lsm_grid_polygon.bounds)

    out_weight_table_file = os.path.join(OUTPUT_DIR,
                                         'weight_table_test_gldas2_1.csv')

    if os.path.exists(out_weight_table_file):
        os.remove(out_weight_table_file)

    weight.write_weight_table(catchment_layer, out_weight_table_file,
                              connect_rivid_list,
                              lsm_grid_lat, lsm_grid_lon,
                              catchment_rivid_list, lsm_grid_rtree,
                              lsm_grid_voronoi_feature_list)

    entry = np.genfromtxt(out_weight_table_file, skip_header=1,
                          delimiter=',')

    expected = np.array([17880832, 7220569.5496, 294, 393, 1, -106.3750,
                         38.3750, 1788083203930294])

    assert_array_equal(entry, expected)

def test_generate_weight_table():
    """
    Verify that the main routine in weight.py produces a weight table
    consistent with benchmark.
    """
    catchment_shapefile = os.path.join(DATA_DIR, 'catchment',
                                       'saguache_colorado', 'catchment.shp')
    connectivity_file = os.path.join(DATA_DIR, 'connectivity',
                                     'rapid_connect_saguache.csv')
    lsm_file = os.path.join(DATA_DIR, 'lsm_grids', 'gldas2',
                            'GLDAS_NOAH025_3H.A20101231.0000.020.nc4')
    out_weight_table_file = os.path.join(OUTPUT_DIR,
                                         'weight_table_test_gldas2_2.csv')
    benchmark_file = os.path.join(DATA_DIR, 'weight_table',
                                  'weight_gldas2.csv')

    weight.generate_weight_table(lsm_file, catchment_shapefile,
                                 connectivity_file, out_weight_table_file)

    compare_weight_table_files(out_weight_table_file, benchmark_file)

def test_generate_weight_table_land_fraction_array():
    """
    Verify that the main routine in weight.py produces a weight table
    consistent with benchmark using a land-sea mask.
    """
    catchment_shapefile = os.path.join(DATA_DIR, 'catchment',
                                       'mendocino_nhdplus',
                                       'NHDCat_mendocino_watershed_hopland_sample.shp')
    connectivity_file = os.path.join(DATA_DIR, 'connectivity',
                                     'rapid_connectivity_mendocino_sample.csv')
    lsm_file = os.path.join(DATA_DIR, 'lsm_grids', 'era5_land_mask',
                            'era5_land-sea_mask_mendocino_subset.nc')
    lsm_lat_variable = 'latitude'
    lsm_lon_variable = 'longitude'
    out_weight_table_file = os.path.join(OUTPUT_DIR,
                                         'weight_table_test_era5_land_mask.csv')
    benchmark_file = os.path.join(DATA_DIR, 'weight_table',
                                  'weight_era5_land_mask.csv')
    lsm_longitude_shift = 1
    lsm_land_fraction_var = 'lsm'

    weight.generate_weight_table(lsm_file, catchment_shapefile,
                                 connectivity_file,
                                 out_weight_table_file,
                                 lsm_lat_variable=lsm_lat_variable,
                                 lsm_lon_variable=lsm_lon_variable,
                                 lsm_longitude_shift=lsm_longitude_shift,
                                 lsm_land_fraction_var=lsm_land_fraction_var)

    compare_weight_table_files(out_weight_table_file, benchmark_file)

def test_generate_weight_table_missing_intersections():
    """
    Verify that the main routine in weight.py produces a weight table
    consistent with benchmark using a land-sea mask.
    """
    catchment_shapefile = os.path.join(DATA_DIR, 'catchment',
                                       'thames_near_cirencester_uk',
                                       'Catchment_thames_drainID45390.shp')
    connectivity_file = os.path.join(DATA_DIR, 'connectivity',
                                     'rapid_connect_45390.csv')
    lsm_file = os.path.join(DATA_DIR, 'lsm_grids', 'lis',
                            'LIS_HIST_201101210000.d01.nc')
    lsm_lat_variable = 'lat'
    lsm_lon_variable = 'lon'
    out_weight_table_file = os.path.join(OUTPUT_DIR,
                                    'weight_table_test_lis_no_intersect.csv')
    benchmark_file = os.path.join(DATA_DIR, 'weight_table',
                                  'weight_lis_no_intersect.csv')

    weight.generate_weight_table(lsm_file, catchment_shapefile,
                                 connectivity_file,
                                 out_weight_table_file,
                                 lsm_lat_variable=lsm_lat_variable,
                                 lsm_lon_variable=lsm_lon_variable,
                                 catchment_id_field_name='DrainLnID')

    compare_weight_table_files(out_weight_table_file, benchmark_file)


#程序文件Pex20_20.py
import glob, numpy as np
from PIL import Image
f = glob.glob("附件1\\*.bmp")  #读入附件1下所有bmp文件名称
n = len(f); a = np.array(Image.open(f[0]))
a = a.astype(float); jj = np.arange(1,n)
L1 = a[:,0]; L2 = a[:,-1]  #拼接的大图像左边界和右边界初始化

tind = [0]
cont_img = a
for i in range(n - 1):
    tcha1 = [];     tcha2 = []
    for j in jj:
        a2 = np.array(Image.open(f[j])).astype(float)
        e1 = a2[:, 0]; e2 = a2[:, -1]
        cha1 = abs(L1 - e2).sum(); cha2 = abs(L2 - e1).sum()
        tcha1.append(cha1) ; tcha2.append(cha2)  #左右边界的差异
    m1 = np.array(tcha1).min(); m2 = np.array(tcha2).min()

    if abs(L1 - 255.).sum() < 1:
        ind = np.where(tcha2 == m2)
        tind=np.hstack((tind, jj[ind]))  # 右拼接
        tt = np.array(Image.open(f[jj[ind[0][0]]])).astype(float)
        cont_img = np.hstack((cont_img, tt))
        L2 = tt[:, -1]; np.delete(jj, ind)
    elif abs(L2 - 255.).sum() < 1:
        ind = np.where(tcha1 == m1)
        tind = np.hstack((jj[ind], tind)) # 左拼接
        tt = np.array(Image.open(f[jj[ind[0][0]]])).astype(float)
        cont_img = np.hstack((tt,cont_img))
        L1 = tt[:, 0]; np.delete(jj, ind)
    elif m1 < m2:
        ind = np.where(tcha1 == m1)
        tind = np.hstack((jj[ind], tind)) # 左拼接
        tt = np.array(Image.open(f[jj[ind[0][0]]])).astype(float)
        cont_img = np.hstack((tt, cont_img))
        L1 = tt[:, 0]; np.delete(jj, ind)
    else:
        ind = np.where(tcha2 == m2)
        tind=np.hstack((tind, jj[ind]))  # 右拼接
        tt = np.array(Image.open(f[jj[ind[0][0]]])).astype(float)
        cont_img = np.hstack((cont_img, tt))
        L2 = tt[:, -1]; np.delete(jj, ind)
print(tind)  #显示各个图像的拼接排列次序
Image.fromarray(cont_img).show()


import pandas as pd
import matplotlib.pyplot as plt

excel_data_df = pd.read_csv('csvatten.csv')
V = excel_data_df['Auth'].tolist()
df = pd.DataFrame({'Auth': excel_data_df['Auth'].tolist()})
counts = df['Auth'].value_counts()
counts.plot.pie(autopct='%.2f%%')
plt.show()

# import pandas as pd
# import matplotlib.pyplot as plt
# csv_file='csvatten.csv'
# data = pd.read_csv(csv_file)
# Name = data["Name"]
# Auth = data["Auth"]
# x=[]
# y=[]
# x=list(Auth)
# y=list(Name)
# plt.scatter(x,y)
# plt.xlabel('Auth->')
# plt.ylabel('Name->')
# plt.title('Data')
# plt.show()
# # plt.pie(x,labels=y,autopct='%.2f%%')


import sys, os, platform
os_name=platform.system()
if (os_name=='Darwin'):
    rootpath='/Users/weiwilliam'
    rootarch='/Volumes/WD2TB/ResearchData'
elif (os_name=='Windows'):
    rootpath='F:\GoogleDrive_NCU\Albany'
    rootarch='F:\ResearchData'
    rootgit='F:\GitHub\swei_research'
elif (os_name=='Linux'):
    if (os.path.exists('/scratch1')):
        rootpath='/scratch2/BMC/gsd-fv3-dev/Shih-wei.Wei'
        rootarch='/scratch2/BMC/gsd-fv3-dev/Shih-wei.Wei/ResearchData'
        rootgit='/home/Shih-wei.Wei/research'
    elif (os.path.exists('/glade')):
        rootpath='/glade/work/swei/output/images'
        rootarch='/scratch2/BMC/gsd-fv3-dev/Shih-wei.Wei/ResearchData'
        rootgit='/glade/u/home/swei/research'
        machine='Cheyenne'
    elif (os.path.exists('/cardinal')):
        rootpath='/data/users/swei/Images'
        rootarch='/scratch2/BMC/gsd-fv3-dev/Shih-wei.Wei/ResearchData'
        rootgit='/home/swei/research'
        machine='S4'
sys.path.append(rootgit+'/pyscripts/functions')
from utils import setup_cmap, ndate
from plot_utils import setupax_2dmap,set_size
import setuparea as setarea
import xarray as xa
import numpy as np
import pandas as pd
from datetime import datetime, timedelta
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.colors as mpcrs
import cartopy.crs as ccrs

# Plotting setup
mpl.rc('axes',titlesize=18,labelsize=16)
mpl.rc('xtick',labelsize=16)
mpl.rc('ytick',labelsize=16)
mpl.rc('legend',fontsize='xx-large')
axe_w=10; axe_h=5
quality=300

if (machine=='Cheyenne'):
   inputpath='/glade/work/dfgrogan/UFS/WM_DTAER/AER'
elif (machine=='S4'):
   inputpath='/data/users/swei/common/MERRA2'

outputpath=rootpath+'/Dataset/MERRA-2/2dmap/CMass_frac'
if ( not os.path.exists(outputpath) ):
    os.makedirs(outputpath)

sdate=2020062212
edate=2020062212
hint=6
area='Glb'
m2tag='inst3_3d_aer_Nv'
tkfreq=2
#species_lst=['dust','seas']
species_lst=['dust','seas','carbon','sulf']

#
clridx=np.array((2,8,9,10,12,14,15,16,18))
cn_cmap=setup_cmap('precip3_16lev',clridx)
cnlvs=np.array((0., 0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.8, 1.0))
#clridx=[0,11,29,38,56,74,83,92,119,128]
#cn_cmap=setup_cmap('MPL_YlOrBr',clridx)
#cnlvs=[0., 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
norm = mpcrs.BoundaryNorm(cnlvs,len(cnlvs))

# Constant configuration
proj=ccrs.PlateCarree()
grav=9.80665e+0

# Setup plotting area
minlon, maxlon, minlat, maxlat, crosszero, cyclic=setarea.setarea(area)
print(minlat,maxlat,minlon,maxlon,crosszero,cyclic)
if (area=='Glb'):
   minlon=-180. ; maxlon=180.
cornerll=[minlat,maxlat,minlon,maxlon]

dust_lst=['DU001','DU002','DU003','DU004','DU005']
seas_lst=['SS001','SS002','SS003','SS004','SS005']
carbon_lst=['OCPHILIC','OCPHOBIC','BCPHILIC','BCPHOBIC']
sulf_lst=['SO4']

syy=int(str(sdate)[:4]); smm=int(str(sdate)[4:6])
sdd=int(str(sdate)[6:8]); shh=int(str(sdate)[8:10])
eyy=int(str(edate)[:4]); emm=int(str(edate)[4:6])
edd=int(str(edate)[6:8]); ehh=int(str(edate)[8:10])

date1 = datetime(syy,smm,sdd,shh)
date2 = datetime(eyy,emm,edd,ehh)
delta = timedelta(hours=hint)
dates = pd.date_range(start=date1, end=date2, freq=delta)

tnum=0
dlist=[]
cdate=sdate
while (cdate<=edate):
    dlist.append(str(cdate))
    tnum=tnum+1
    cdate=ndate(hint,cdate)

dates_count=0
for date in dlist:
    yy=date[:4] ; mm=date[4:6] ; dd=date[6:8] ; hh=date[8:10]
    pdy=date[:8]
    if (yy=='2020' and mm=='09'):
       m2ind='401'
    else:
       m2ind='400'
# MERRA2_401.inst3_3d_aer_Nv.20200916_12Z.nc4
    if (machine=='Cheyenne'):
       infile=inputpath+'/MERRA2_'+m2ind+'.'+m2tag+'.'+pdy+'_'+hh+'Z.nc4'
       multi_time=0
    if (machine=='S4'):
       infile=inputpath+'/'+yy+'/'+mm+'/MERRA2_'+m2ind+'.'+m2tag+'.'+pdy+'.nc4'
       multi_time=1

    ds=xa.open_dataset(infile)
    if (multi_time):
       ds=ds.sel(time=dates[dates_count])
    else:
       ds=ds.sel(time=ds.time[0])
    if (area!='Glb'):
       ds=ds.sel(lon=slice(minlon,maxlon),lat=slice(minlat,maxlat))

    delp=ds.DELP
    kgkg_kgm2=delp/grav

    spidx=0
    for sp in species_lst:
        spc_conc=0.
        if (sp=='dust'):
           varlst=dust_lst
        if (sp=='seas'):
           varlst=seas_lst
        if (sp=='carbon'):
           varlst=carbon_lst
        if (sp=='sulf'):
           varlst=sulf_lst
        varidx=0
        for var in varlst:
           tmp=ds[var]*kgkg_kgm2
           if (sp=='dust'):
              if (varidx==0):
                 dust_conc=tmp
              else:
                 dust_conc=xa.concat((dust_conc,tmp),dim='bins')
           if (sp=='seas'):
              if (varidx==0):
                 seas_conc=tmp
              else:
                 seas_conc=xa.concat((seas_conc,tmp),dim='bins')
           if (sp=='carbon'):
              if (varidx==0):
                 carbon_conc=tmp
              else:
                 carbon_conc=xa.concat((carbon_conc,tmp),dim='bins')
           if (sp=='sulf'):
              if (varidx==0):
                 sulf_conc=tmp
              else:
                 sulf_conc=xa.concat((sulf_conc,tmp),dim='bins')
           spc_conc=spc_conc+tmp
           varidx+=1

        if (spidx==0):
           conc=spc_conc
        else:
           conc=xa.concat((conc,spc_conc),dim='species')
        spidx+=1

# Get the column integral
    cmass=conc.sum(dim='lev')
    cmass_total=cmass.sum(dim='species')

    spidx=0
    for sp in species_lst:
        cblb='%s Fraction of column mass density' %(sp)
        title='%s ' %(species_lst[spidx])
        outname='%s/%s_%s_cmass_frac.%s.png'  %(outputpath,area,species_lst[spidx],date)

        pltdata=cmass[spidx,:,:]/cmass_total
        fig,ax=setupax_2dmap(cornerll,area,proj,lbsize=16.)
        set_size(axe_w,axe_h,b=0.13,l=0.05,r=0.95,t=0.95)
        cn=ax.contourf(pltdata.lon,pltdata.lat,pltdata,levels=cnlvs,cmap=cn_cmap,norm=norm)
        ax.set_title(title)
        plt.colorbar(cn,ax=ax,orientation='horizontal',ticks=cnlvs[::tkfreq],
                     format='%.2f',fraction=0.045,aspect=40,pad=0.08,label=cblb)
        print(outname)
        fig.savefig(outname,dpi=quality)
        plt.close()

        if (sp=='dust'):
           varlst=dust_lst
           sp_cmass=dust_conc.sum(dim='lev')
        elif (sp=='seas'):
           varlst=seas_lst
           sp_cmass=seas_conc.sum(dim='lev')
        elif (sp=='carbon'):
           varlst=carbon_lst
           sp_cmass=carbon_conc.sum(dim='lev')
        elif (sp=='sulf'):
           varlst=sulf_lst
           sp_cmass=sulf_conc.sum(dim='lev')
        subspidx=0
        for subsp in varlst:
            cblb='%s Fraction of column mass density' %(subsp)
            title='%s ' %(subsp)
            outname='%s/%s_%s_cmass_frac.%s.png'  %(outputpath,area,subsp,date)
            if (sp=='sulf'):
               pltdata=sp_cmass/cmass_total
            else:
               pltdata=sp_cmass[subspidx,:,:]/cmass_total
            fig,ax=setupax_2dmap(cornerll,area,proj,lbsize=16.)
            set_size(axe_w,axe_h,b=0.13,l=0.05,r=0.95,t=0.95)
            cn=ax.contourf(pltdata.lon,pltdata.lat,pltdata,levels=cnlvs,cmap=cn_cmap,norm=norm)
            ax.set_title(title)
            plt.colorbar(cn,ax=ax,orientation='horizontal',ticks=cnlvs[::tkfreq],
                         format='%.2f',fraction=0.045,aspect=40,pad=0.08,label=cblb)
            print(outname)
            fig.savefig(outname,dpi=quality)
            plt.close()
            subspidx+=1

        spidx+=1

    dates_count+=1