from cube.data.objects import Doc, Sentence, Word, Token

class ComponentType:
    Tokenizer = 1
    CWExpander = 2
    POSTagger = 3
    Lemmatizer = 4
    Parser = 5
    NER = 6

class Model:
    def __init__(self):
        pass

    def __call__(self, input_object, **kwargs):
        return None

class Component():
    def __init__(self, use_gpu: bool = True, gpu_batch_size: int = 1):
        self.input_format = Doc  # type of data object or str
        self.output_format = Doc  # type of data object
        self.depends = []  # list of other components
        self.provides = []  # list of other components

        self.model_filepath = None

        self.use_gpu = use_gpu
        self.gpu_batch_size = gpu_batch_size

        self.model = None

    def load_model(self, model_path):
        pass

    def process(self, input_object):
        assert (self.model is not None), "Model is none, please load model first"
        return self.model(input_object=input_object)

class TokenizerComponent(Component):
    def __init__(self):
        super().__init__()
        self.input_format = str
        self.depends = []
        self.provides = [ComponentType.Tokenizer]

class CWExpanderComponent(Component):
    def __init__(self):
        super().__init__()
        self.depends = [ComponentType.Tokenizer]
        self.provides = [ComponentType.CWExpander]

class POSTaggerComponent(Component):
    def __init__(self):
        super().__init__()
        self.depends = [ComponentType.Tokenizer, ComponentType.CWExpander]
        self.provides = [ComponentType.POSTagger]

class LemmatizerComponent(Component):
    def __init__(self):
        super().__init__()
        self.depends = [ComponentType.Tokenizer, ComponentType.CWExpander, ComponentType.Parser]
        self.provides = [ComponentType.Lemmatizer]

class ParserComponent(Component):
    def __init__(self):
        super().__init__()
        self.depends = [ComponentType.Tokenizer, ComponentType.CWExpander]
        self.provides = [ComponentType.Parser]

class NERComponent(Component):
    def __init__(self):
        super().__init__()
        self.depends = [ComponentType.Tokenizer, ComponentType.CWExpander, ComponentType.Parser]
        self.provides = [ComponentType.NER]

class Pipeline():
    def is_valid(components: Component):
        available = set()
        required = set()
        for component in components:
            available |= set(component.provides)
            required |= set(component.depends)
        return required.issubset(available)


class DFS:
	def __init__(self, size):
		self.size = size
		self.color = ['white']*size
		# self.graph = {0:[1], 1:[0,2], 2:[1,3], 3:[2,4,5], 4:[3,5,7], 5:[3,4,6,7], 6:[5,7], 7:[4,5,6]}
		self.graph = {0:[1,2], 1:[2], 2:[3], 3:[1], 4:[3,5], 5:[5]}

	def dfsTraversal(self):
		for i in range(self.size):
			if self.color[i] == 'white':
				self.dfs(i)

	def dfs(self, source):
		stack = []
		stack.append(source)
		self.color[source] = 'gray'

		while stack:
			current = stack.pop()
			print(current, end=' ')
			for i in range(len(self.graph[current])):
				if self.color[self.graph[current][i]] == 'white':
					self.color[self.graph[current][i]] = 'gray'
					stack.append(self.graph[current][i])
			self.color[current] = 'black'

dfs = DFS(6)
dfs.dfsTraversal()


import json
import uuid
from copy import deepcopy

import pytest
from pendulum import DateTime, datetime

from gsy_framework.data_classes import (
    BidOfferMatch, BaseBidOffer, Offer, Bid, json_datetime_serializer,
    TradeBidOfferInfo, Trade, BalancingOffer, BalancingTrade, Clearing,
    MarketClearingState, TraderDetails)
from gsy_framework.utils import datetime_to_string_incl_seconds, limit_float_precision

DEFAULT_DATETIME = datetime(2021, 11, 3, 10, 45)

def test_json_datetime_serializer():
    """Test if json_datetime_serializer works as expected."""
    assert json_datetime_serializer("") is None

    my_date_time = DateTime.today()
    assert json_datetime_serializer(my_date_time) == datetime_to_string_incl_seconds(my_date_time)

class TestBidOfferMatch:
    """Tester class for the BidOfferMatch dataclass."""

    @staticmethod
    def test_serializable_dict():
        """Test the serializable_dict method of BidOfferMatch dataclass."""
        bid_offer_match = BidOfferMatch(
            market_id="market_id",
            time_slot="2021-10-06T12:00",
            bid={"type": "bid"},
            offer={"type": "offer"},
            selected_energy=1,
            trade_rate=1)
        expected_dict = {"market_id": "market_id",
                         "time_slot": "2021-10-06T12:00",
                         "bid": {"type": "bid"},
                         "offer": {"type": "offer"},
                         "selected_energy": 1,
                         "trade_rate": 1,
                         "matching_requirements": None}
        assert bid_offer_match.serializable_dict() == expected_dict

    @staticmethod
    def test_is_valid_dict():
        """Test the is_valid_dict method of BidOfferMatch dataclass."""
        bid_offer_match = {"market_id": "market_id",
                           "time_slot": "2021-10-06T12:00",
                           "bid": {"type": "bid"},
                           "offer": {"type": "offer"},
                           "selected_energy": 1,
                           "trade_rate": 1}
        assert BidOfferMatch.is_valid_dict(bid_offer_match)

        # Key does not exist
        bid_offer_match = {"market_id": "market_id",
                           "time_slot": "2021-10-06T12:00",
                           "bid": {"type": "bid"},
                           "offer": {"type": "offer"},
                           "selected_energy": 1,
                           }
        assert not BidOfferMatch.is_valid_dict(bid_offer_match)

        # Wrong type
        bid_offer_match = {"market_id": "market_id",
                           "time_slot": "2021-10-06T12:00",
                           "bid": {"type": "bid"},
                           "offer": {"type": "offer"},
                           "selected_energy": 1,
                           "trade_rate": ""}
        assert not BidOfferMatch.is_valid_dict(bid_offer_match)

        # No market id
        bid_offer_match = {"market_id": None,
                           "time_slot": "2021-10-06T12:00",
                           "bid": {"type": "bid"},
                           "offer": {"type": "offer"},
                           "selected_energy": 1,
                           "trade_rate": ""}
        assert not BidOfferMatch.is_valid_dict(bid_offer_match)

        # No selected energy
        bid_offer_match = {"market_id": "market_id",
                           "time_slot": "2021-10-06T12:00",
                           "bid": {"type": "bid"},
                           "offer": {"type": "offer"},
                           "selected_energy": None,
                           "trade_rate": ""}
        assert not BidOfferMatch.is_valid_dict(bid_offer_match)

    @staticmethod
    def test_from_dict():
        """Test the from_dict method of BidOfferMatch dataclass."""
        expected_dict = {"market_id": "market_id",
                         "time_slot": "2021-10-06T12:00",
                         "bid": {"type": "bid"},
                         "offer": {"type": "offer"},
                         "selected_energy": 1,
                         "trade_rate": 1}
        bid_offer_match = BidOfferMatch.from_dict(expected_dict)
        assert bid_offer_match.market_id == expected_dict["market_id"]
        assert bid_offer_match.time_slot == expected_dict["time_slot"]
        assert bid_offer_match.bid == expected_dict["bid"]
        assert bid_offer_match.offer == expected_dict["offer"]
        assert bid_offer_match.selected_energy == expected_dict["selected_energy"]
        assert bid_offer_match.trade_rate == expected_dict["trade_rate"]

    @staticmethod
    def test_from_dict_returns_none_for_invalid_dict():
        assert BidOfferMatch.from_dict({}) is None

    @staticmethod
    def test_bid_energy():
        """Test whether the bid_energy property returns the correct value."""
        orders_match = BidOfferMatch(bid={"energy": 10}, offer={}, market_id="",
                                     time_slot="", selected_energy=1, trade_rate=1)
        assert orders_match.bid_energy == 10

        # When there is a bid_requirement, it should be prioritized
        orders_match.matching_requirements = {"bid_requirement": {"energy": 5}}
        assert orders_match.bid_energy == 5

    @staticmethod
    def test_bid_energy_rate():
        orders_match = BidOfferMatch(bid={"energy_rate": 2,
                                          "energy": 1}, offer={}, market_id="",
                                     time_slot="", selected_energy=1, trade_rate=1)
        assert orders_match.bid_energy_rate == 2
        # When there is a bid_requirement, it should be prioritized
        orders_match.matching_requirements = {"bid_requirement": {"price": 4}}
        assert orders_match.bid_energy_rate == (
            orders_match.matching_requirements["bid_requirement"]["price"]
            / orders_match.selected_energy)

class TestBaseBidOffer:
    """Test BaseBidOffer class."""

    def setup_method(self):
        self._seller = TraderDetails("seller", str(uuid.uuid4()))
        self._buyer = TraderDetails("buyer", str(uuid.uuid4()))
        self.initial_data = {
            "id": str(uuid.uuid4()),
            "creation_time": DEFAULT_DATETIME,
            "time_slot": DEFAULT_DATETIME,
            "price": 10,
            "energy": 30,
            "original_price": 8,
        }

    def test_init(self):
        """Test __init__."""
        bid_offer = BaseBidOffer(
            **self.initial_data
        )
        assert bid_offer.id == str(self.initial_data["id"])
        assert bid_offer.creation_time == self.initial_data["creation_time"]
        assert bid_offer.price == self.initial_data["price"]
        assert bid_offer.energy == self.initial_data["energy"]
        assert bid_offer.original_price == self.initial_data["original_price"]
        assert bid_offer.energy_rate == limit_float_precision(self.initial_data["price"] /
                                                              self.initial_data["energy"])

        # Test whether the original_price will resort to the "price" member
        self.initial_data.pop("original_price")
        bid_offer = BaseBidOffer(
            **self.initial_data
        )
        assert bid_offer.original_price == self.initial_data["price"]

    def test_update_price(self):
        bid_offer = BaseBidOffer(
            **self.initial_data
        )
        bid_offer.update_price(30)
        assert bid_offer.price == 30
        assert bid_offer.energy_rate == 30 / bid_offer.energy

    def test_update_energy(self):
        bid_offer = BaseBidOffer(
            **self.initial_data
        )
        bid_offer.update_energy(40)
        assert bid_offer.energy == 40
        assert bid_offer.energy_rate == bid_offer.price / 40

    def test_to_json_string(self):
        bid_offer_keys = {
            "id", "creation_time", "time_slot", "original_price", "price", "energy",
            "type", "energy_rate"}
        bid_offer = BaseBidOffer(
            **self.initial_data
        )
        obj_dict = json.loads(bid_offer.to_json_string())
        assert set(obj_dict.keys()) == bid_offer_keys

        assert json.dumps(obj_dict, sort_keys=True) == json.dumps(
            {key: getattr(bid_offer, key) for key in bid_offer_keys}, sort_keys=True,
            default=json_datetime_serializer
        )

    def test_serializable_dict(self):
        bid_offer = BaseBidOffer(
            **self.initial_data
        )
        assert bid_offer.serializable_dict() == {
            "type": "BaseBidOffer",
            "id": str(bid_offer.id),
            "energy": bid_offer.energy,
            "price": bid_offer.price,
            "energy_rate": bid_offer.energy_rate,
            "original_price": bid_offer.original_price,
            "creation_time": datetime_to_string_incl_seconds(bid_offer.creation_time),
            "time_slot": datetime_to_string_incl_seconds(bid_offer.time_slot)
        }

    def test_from_json(self):
        offer = Offer(
            **self.initial_data,
            seller=self._seller
        )
        offer_json = offer.to_json_string()
        assert offer == Offer.from_json(offer_json)

        bid = Bid(
            **self.initial_data,
            buyer=self._buyer
        )
        bid_json = bid.to_json_string()
        assert bid == Bid.from_json(bid_json)

    @pytest.mark.parametrize("time_stamp", [None, DEFAULT_DATETIME])
    def test_from_json_deals_with_time_stamps_correctly(self, time_stamp):
        updated_initial_data = self.initial_data
        updated_initial_data.update({"time_slot": time_stamp,
                                     "creation_time": time_stamp})
        bid = Bid(**updated_initial_data, buyer=self._buyer)
        bid_json = bid.to_json_string()
        bid = Bid.from_json(bid_json)
        assert bid.creation_time == time_stamp

    def test_from_json_asserts_if_no_type_was_provided(self):
        bid = Bid(**self.initial_data, buyer=self._buyer)
        bid_json = bid.to_json_string()
        bid_json = bid_json.replace('"type": "Bid", ', "")
        with pytest.raises(AssertionError):
            Bid.from_json(bid_json)

    def test_from_json_asserts_if_wrong_type_was_provided(self):
        bid = Bid(**self.initial_data, buyer=self._buyer)
        bid_json = bid.to_json_string()
        bid_json = bid_json.replace('"type": "Bid", ', '"type": "InvalidType", ')
        with pytest.raises(AssertionError):
            Bid.from_json(bid_json)

class TestOffer:
    def setup_method(self):
        self.initial_data = {
            "id": uuid.uuid4(),
            "creation_time": DEFAULT_DATETIME,
            "price": 10,
            "energy": 30,
            "original_price": 8,
            "seller": TraderDetails("seller", str(uuid.uuid4())),
            "time_slot": DEFAULT_DATETIME
        }

    def test_init(self):
        offer = Offer(
            **self.initial_data
        )
        assert offer.id == str(self.initial_data["id"])
        assert offer.creation_time == self.initial_data["creation_time"]
        assert offer.time_slot == self.initial_data["time_slot"]
        assert offer.price == self.initial_data["price"]
        assert offer.energy == self.initial_data["energy"]
        assert offer.original_price == self.initial_data["original_price"]
        assert offer.energy_rate == limit_float_precision(self.initial_data["price"] /
                                                          self.initial_data["energy"])
        assert offer.seller.name == "seller"
        assert offer.seller.uuid == self.initial_data["seller"].uuid
        assert offer.seller.origin is None
        assert offer.seller.origin_uuid is None

    def test_hash(self):
        offer = Offer(
            **self.initial_data,
        )
        assert offer.__hash__() == hash(offer.id)

    def test_repr(self):
        offer = Offer(
            **self.initial_data,
        )
        assert (repr(offer) ==
                f"<Offer('{offer.id!s:.6s}', '{offer.energy} kWh@{offer.price}',"
                f" '{offer.seller.name} {offer.energy_rate}'>")

    def test_str(self):
        offer = Offer(
            **self.initial_data,
        )
        assert (str(offer) ==
                f"{{{offer.id!s:.6s}}} [origin: {offer.seller.origin}] "
                f"[{offer.seller.name}]: {offer.energy} kWh @ {offer.price} @ {offer.energy_rate}")

    def test_serializable_dict(self):
        offer = Offer(
            **self.initial_data,
        )
        assert offer.serializable_dict() == {
            "type": "Offer",
            "id": str(offer.id),
            "energy": offer.energy,
            "price": offer.price,
            "energy_rate": offer.energy_rate,
            "original_price": offer.original_price,
            "creation_time": datetime_to_string_incl_seconds(offer.creation_time),
            "seller": offer.seller.serializable_dict(),
            "time_slot": datetime_to_string_incl_seconds(offer.time_slot)
        }

    def test_from_dict(self):
        offer = Offer(
            **self.initial_data,
        )
        assert Offer.from_dict(offer.serializable_dict()) == offer

    def test_eq(self):
        offer = Offer(
            **self.initial_data
        )
        other_offer = Offer(
            **self.initial_data
        )
        assert offer == other_offer

        other_offer.id = "other_id"
        assert offer != other_offer

    def test_csv_values(self):
        offer = Offer(
            **self.initial_data
        )
        rate = round(offer.energy_rate, 4)
        assert offer.csv_values() == (
            offer.creation_time, rate, offer.energy, offer.price, offer.seller.name,
            offer.seller.origin)

    @staticmethod
    def test_csv_fields():
        assert (Offer.csv_fields() ==
                ("creation_time", "rate [ct./kWh]", "energy [kWh]", "price [ct.]", "seller",
                 "seller origin"))

    def test_copy(self):
        offer = Offer(
            **self.initial_data
        )
        second_offer = Offer.copy(offer)
        assert offer == second_offer

    def test_accumulated_grid_fees(self):
        offer = Offer(**self.initial_data)
        assert offer.accumulated_grid_fees == offer.price - offer.original_price

class TestBid:
    def setup_method(self):
        self.initial_data = {
            "id": uuid.uuid4(),
            "creation_time": DEFAULT_DATETIME,
            "price": 10,
            "energy": 30,
            "original_price": 8,
            "buyer": TraderDetails("buyer", str(uuid.uuid4())),
            "time_slot": DEFAULT_DATETIME
        }

    def test_init(self):
        bid = Bid(
            **self.initial_data
        )
        assert bid.id == str(self.initial_data["id"])
        assert bid.creation_time == self.initial_data["creation_time"]
        assert bid.price == self.initial_data["price"]
        assert bid.energy == self.initial_data["energy"]
        assert bid.original_price == self.initial_data["original_price"]
        assert bid.energy_rate == limit_float_precision(self.initial_data["price"] /
                                                        self.initial_data["energy"])
        assert bid.buyer.name == "buyer"
        assert bid.buyer.uuid == self.initial_data["buyer"].uuid
        assert bid.buyer.origin is None
        assert bid.buyer.origin_uuid is None

    def test_hash(self):
        bid = Bid(
            **self.initial_data
        )
        assert bid.__hash__() == hash(bid.id)

    def test_repr(self):
        bid = Bid(
            **self.initial_data
        )
        assert (repr(bid) ==
                f"<Bid {{{bid.id!s:.6s}}} [{bid.buyer.name}] "
                f"{bid.energy} kWh @ {bid.price} {bid.energy_rate}>")

    def test_str(self):
        bid = Bid(
            **self.initial_data
        )
        assert (str(bid) ==
                f"{{{bid.id!s:.6s}}} [origin: {bid.buyer.origin}] [{bid.buyer.name}] "
                f"{bid.energy} kWh @ {bid.price} {bid.energy_rate}")

    def test_serializable_dict(self):
        bid = Bid(
            **self.initial_data
        )

        assert bid.serializable_dict() == {
            "type": "Bid",
            "id": str(bid.id),
            "energy": bid.energy,
            "price": bid.price,
            "energy_rate": bid.energy_rate,
            "original_price": bid.original_price,
            "creation_time": datetime_to_string_incl_seconds(bid.creation_time),
            "buyer": bid.buyer.serializable_dict(),
            "time_slot": datetime_to_string_incl_seconds(bid.time_slot)
        }

    def test_from_dict(self):
        bid = Bid(
            **self.initial_data,
        )
        assert Bid.from_dict(bid.serializable_dict()) == bid

    def test_eq(self):
        bid = Bid(
            **self.initial_data
        )
        other_bid = Bid(
            **self.initial_data
        )
        assert bid == other_bid

        other_bid.id = "other_id"
        assert bid != other_bid

    def test_csv_values(self):
        bid = Bid(
            **self.initial_data
        )
        rate = round(bid.energy_rate, 4)
        assert bid.csv_values() == (bid.creation_time, rate, bid.energy, bid.price, bid.buyer.name,
                                    bid.buyer.origin)

    @staticmethod
    def test_csv_fields():
        assert (Bid.csv_fields() ==
                ("creation_time", "rate [ct./kWh]", "energy [kWh]", "price [ct.]", "buyer",
                 "buyer origin"))

    def test_accumulated_grid_fees(self):
        bid = Bid(**self.initial_data)
        assert bid.accumulated_grid_fees == bid.original_price - bid.price

class TestTradeBidOfferInfo:
    @staticmethod
    def test_to_json_string():
        trade_bid_offer_info = TradeBidOfferInfo(
            1, 1, 1, 1, 1
        )
        assert (trade_bid_offer_info.to_json_string() ==
                json.dumps(trade_bid_offer_info.serializable_dict()))

    @staticmethod
    def test_from_json():
        trade_bid_offer_info = TradeBidOfferInfo(
            1, 1, 1, 1, 1
        )
        trade_bid_offer_info_json = trade_bid_offer_info.to_json_string()
        assert (TradeBidOfferInfo.from_json(trade_bid_offer_info_json) ==
                trade_bid_offer_info)

class TestTrade:
    def setup_method(self):
        self._seller = TraderDetails("seller", str(uuid.uuid4()))
        self._buyer = TraderDetails("buyer", str(uuid.uuid4()))
        self.initial_data = {
            "id": "my_id",
            "creation_time": DEFAULT_DATETIME,
            "offer": Offer("id", DEFAULT_DATETIME, 1, 2, seller=self._seller),
            "seller": self._seller,
            "buyer": self._buyer,
            "traded_energy": 1,
            "trade_price": 1,
            "matching_requirements": {"requirement": "value"}}

    def test_str(self):
        trade = Trade(**self.initial_data)
        assert (str(trade) ==
                f"{{{trade.id!s:.6s}}} [origin: {trade.seller.origin} -> {trade.buyer.origin}] "
                f"[{trade.seller.name} -> {trade.buyer.name}] {trade.traded_energy} kWh"
                f" @ {trade.trade_price} {round(trade.trade_rate, 8)} "
                f"{trade.match_details['offer'].id}  [fee: {trade.fee_price} cts.] "
                f"{trade.matching_requirements or ''}")

    @staticmethod
    def test_csv_fields():
        assert Trade.csv_fields() == (
            "creation_time", "rate [ct./kWh]", "energy [kWh]", "seller", "seller origin", "buyer",
            "buyer origin")

    def test_csv_values(self):
        trade = Trade(**self.initial_data)
        rate = round(trade.trade_rate, 4)
        assert (trade.csv_values() ==
                (trade.creation_time, rate, trade.traded_energy, trade.seller.name,
                 trade.seller.origin, trade.buyer.name, trade.buyer.origin))

    def test_to_json_string(self):
        trade = Trade(**self.initial_data)
        assert trade.to_json_string() == json.dumps(trade.serializable_dict())

    def test_from_json(self):
        trade = Trade(
            **self.initial_data
        )
        trade.residual = deepcopy(trade.match_details["offer"])
        trade.offer_bid_trade_info = TradeBidOfferInfo(
            1, 1, 1, 1, 1
        )

        trade_json_str = trade.to_json_string()
        assert Trade.from_json(trade_json_str) == trade

    @pytest.mark.parametrize("time_stamp", [None, DEFAULT_DATETIME])
    def test_from_json_deals_with_time_stamps_correctly(self, time_stamp):
        updated_initial_data = self.initial_data
        updated_initial_data.update({"time_slot": time_stamp,
                                     "creation_time": time_stamp})
        trade = Trade(**updated_initial_data)
        trade_json = trade.to_json_string()
        trade = Trade.from_json(trade_json)
        assert trade.creation_time == time_stamp

    def test_is_bid_trade(self):
        trade = Trade(
            **self.initial_data
        )
        assert trade.is_bid_trade is False

        trade.match_details["bid"] = Bid("id", DateTime.now(), 1, 2, self._buyer)
        trade.match_details["offer"] = None
        assert trade.is_bid_trade is True

    def test_is_offer_trade(self):
        trade = Trade(
            **self.initial_data
        )
        assert trade.is_offer_trade is True

        trade.match_details["bid"] = Bid("id", DateTime.now(), 1, 2, self._buyer)
        trade.match_details["offer"] = None
        assert trade.is_offer_trade is False

    def test_serializable_dict(self):
        trade = Trade(
            **{
                "id": "my_id",
                "offer": Offer("id", DEFAULT_DATETIME, 1, 2, self._seller),
                "buyer": TraderDetails(
                    name="buyer",
                    origin="buyer_origin",
                    origin_uuid="buyer_origin_id",
                    uuid="buyer_id",
                ),
                "seller": TraderDetails(
                    name="seller",
                    origin="seller_origin",
                    origin_uuid="seller_origin_id",
                    uuid="seller_id",
                ),
                "traded_energy": 1,
                "trade_price": 1,
                "fee_price": 2,
                "creation_time": DEFAULT_DATETIME,
                "time_slot": DEFAULT_DATETIME
            }
        )
        assert trade.serializable_dict() == {
            "type": "Trade",
            "match_type": "Offer",
            "id": trade.id,
            "offer": trade.match_details["offer"].serializable_dict(),
            "bid": None,
            "residual": trade.residual.serializable_dict() if trade.residual is not None else None,
            "energy": trade.traded_energy,
            "energy_rate": trade.trade_rate,
            "price": trade.trade_price,
            "buyer": {
                "name": trade.buyer.name,
                "uuid": trade.buyer.uuid,
                "origin": trade.buyer.origin,
                "origin_uuid": trade.buyer.origin_uuid
            },
            "seller": {
                "name": trade.seller.name,
                "uuid": trade.seller.uuid,
                "origin": trade.seller.origin,
                "origin_uuid": trade.seller.origin_uuid
            },
            "fee_price": trade.fee_price,
            "creation_time": datetime_to_string_incl_seconds(trade.creation_time),
            "time_slot": datetime_to_string_incl_seconds(trade.creation_time),
            "offer_bid_trade_info": None
        }

class TestBalancingOffer(TestOffer):
    def setup_method(self):
        self.initial_data = {
            "id": uuid.uuid4(),
            "creation_time": DEFAULT_DATETIME,
            "price": 10,
            "energy": 30,
            "original_price": 8,
            "seller": TraderDetails("seller", str(uuid.uuid4())),
            "time_slot": DEFAULT_DATETIME
        }

    def test_repr(self):
        offer = BalancingOffer(**self.initial_data)
        assert (repr(offer) ==
                f"<BalancingOffer('{offer.id!s:.6s}', "
                f"'{offer.energy} kWh@{offer.price}', "
                f"'{offer.seller.name} {offer.energy_rate}'>")

    def test_str(self):
        offer = BalancingOffer(**self.initial_data)
        assert (str(offer) ==
                f"<BalancingOffer{{{offer.id!s:.6s}}} [{offer.seller.name}]: "
                f"{offer.energy} kWh @ {offer.price} @ {offer.energy_rate}>")

class TestBalancingTrade(TestTrade):

    def test_str(self):
        trade = BalancingTrade(**self.initial_data)
        assert (str(trade) ==
                f"{{{trade.id!s:.6s}}} [{trade.seller.name} -> {trade.buyer.name}] "
                f"{trade.traded_energy} kWh @ {trade.trade_price}"
                f" {trade.trade_rate} {trade.match_details['offer'].id}  ")

class TestClearing:
    @staticmethod
    def test_serializable_dict():
        clearing = Clearing(energy=1, rate=1)
        assert clearing.serializable_dict() == {
            "energy": 1, "rate": 1}

class TestMarketClearingState:
    @staticmethod
    def test_csv_fields():
        assert MarketClearingState.csv_fields() == ("creation_time", "rate [ct./kWh]")


import os
import argparse

import matplotlib.pyplot as plt
from datetime import datetime, timedelta, date
import numpy as np
import pints
import pints.io
import pints.plot
from nottingham_covid_modelling import MODULE_DIR
# Load project modules
from nottingham_covid_modelling.lib._command_line_args import NOISE_MODEL_MAPPING
from nottingham_covid_modelling.lib.equations import  get_model_SIUR_solution, get_model_solution, get_model_SIR_solution, get_model_SEIUR_solution
from nottingham_covid_modelling.lib.settings import Params, get_file_name_suffix, get_file_name_suffix_anymodel, get_parameter_names, parameter_to_optimise_list

MODEL_FUNCTIONS ={'SItD':get_model_solution, 'SIR': get_model_SIR_solution, 'SIRDeltaD': get_model_SIR_solution, 'SIUR':get_model_SIUR_solution, \
'SEIUR':get_model_SEIUR_solution}

def plot_mcmc():
    parser = argparse.ArgumentParser()

    parser.add_argument("--show_plots", action='store_true', help="whether to show plots or not", default=False)
    parser.add_argument("--model_name", type=str, help="which model to use", choices=MODEL_FUNCTIONS.keys(), default='SIR')
    parser.add_argument("--out_mcmc", type=str, help="folder to store mcmc runs files in. This folder will be located ./folder, default: ./out-mcmc",
                        default= 'out-mcmc')
    parser.add_argument("-partial", "--fit_partial", action='store_true', help='Whether to fit a subset of the model  parameters (see \"-pto\"), ', default=False)
    parser.add_argument("-pto", "--params_to_optimise", nargs='+', type=str,  \
                        help="If \"--fit_full\" is not given, select which parameters to optimise, e.g. -pto rho Iinit1 eta. This \
                        flag do not include the step parameters (see \"-fitstep\").\
                        OPTIONS FOR EACH MODEL ------ (1)SIR: rho Iinit1 theta; \
                        (2)SIRDeltaD: rho Iinit1 theta DeltaD; (3)SIUR: rho Iinit1 theta xi;  \
                        (4)SEIUR: rho Iinit1 theta eta xi; (5)SItD: rho Iinit1")
    parser.add_argument("--fixed_eta", type=float, help="value of eta. If eta will be fitted this value is ignored. Default value is the best fit for SEIUR model to clinical data: 1/2.333971", default = 0.4285)#0.1923,
    parser.add_argument("--fixed_theta", type=float, help="value of theta. If theta in fitted params, this value is ignored.  Default value is the best fit for SEIUR model to clinical data: 1/3.839233", default=0.2605)
    parser.add_argument("-fitstep", "--fit_step", action='store_false', help='Whether to fit step parameters', default=True)
    parser.add_argument("--syndata_num", type=int, help="Give the number of the synthetic data set you want to fit, default 1", default=1)
    parser.add_argument("--burn_in", help="number of MCMC iterations to ignore",
                        default=25000, type=int)
    parser.add_argument("--chain", type=int, help="which chain to use", default=1)
    parser.add_argument("--informative_priors", action='store_true', help='Whether to use informative priors', default=False)
    parser.add_argument("--chains5", action='store_true', help='Whether to do 5 instead of 3 chains', default=False)

    # At the moment, syntethic data sets 2-9 have travel and step options only. There is only one data sets without step and one with neither travel nor step.

    args = parser.parse_args()
    FitFull = not args.fit_partial
    params_fromOptions = args.params_to_optimise
    FitStep = args.fit_step
    ModelName = args.model_name
    SyntDataNum_file = args.syndata_num
    Fixed_eta = args.fixed_eta
    Fixed_theta = args.fixed_theta

    if FitFull:
        print("Fitting full parameters, any subset of parameters will be ignored. \nIf you want to fit only some parameters change -partial and list -pto to fit ")
    else:
        if params_fromOptions is None:
            parser.error("If -partial, -pto is required. You did not specify -pto.")

    fixed_params_tag = ""
    if not FitFull and ModelName == 'SEIUR':
        if 'eta' not in  params_fromOptions:
            fixed_params_tag = fixed_params_tag + '_fixedEta-' + str(Fixed_eta)
    if not FitFull and ModelName != 'SItD':
        if 'theta' not in  params_fromOptions:
            fixed_params_tag = fixed_params_tag + '_fixedTheta-' + str(Fixed_theta)

    if args.chains5:
        nchains = 5
    else:
        nchains = 3

    # For reproducibility:
    np.random.seed(100)

    # Number of days to fit
    maxtime_fit = 150

    # Get parameters, p
    p = Params()
    # Fixed for the synth data, based on UK google and ONS data:
    p.N = 59.1e6
    p.numeric_max_age = 35
    p.extra_days_to_simulate = 10
    p.IFR = 0.00724 # UK time
    p.square_lockdown = True
    p.alpha = np.ones(p.maxtime)
    p.lockdown_baseline = 0.2814 #0.2042884852266899
    p.lockdown_offset = 31.57 #34.450147247864166
    if not args.informative_priors:
        p.flat_priors = True
    # For saving file names:
    rho_label = '_rho_0-2'
    Noise_label = 'NBphi_2e-3_'

    # Storing parameter values. Default values are given by  best the fit from Kirsty
    if ModelName != 'SItD':
        p.theta = Fixed_theta
        p.eta = Fixed_eta

    if ModelName == 'SEIUR':
        p.xi = 1 / 11.744308
    else:
        p.xi = 1 / (18.69 - 1 - 5.2)

    if ModelName == 'SIRDeltaD':
        p.DeltaD = 18.69 - 1 - 5.2
    else:
        p.DeltaD = 0

    # define the params to optimize
    parameters_to_optimise = parameter_to_optimise_list(FitFull, FitStep, ModelName, params_fromOptions)
    parameter_names = get_parameter_names(parameters_to_optimise, FitStep)

    # Get noise model
    noise_model = NOISE_MODEL_MAPPING['NegBinom']

    # Get simulated Age data from file
    print('Getting simulated data...')

    # folder to load data
    if SyntDataNum_file == 1: # Original default syntethic data
        folder_path =  os.path.join(MODULE_DIR, 'out_SIRvsAGEfits')
        full_fit_data_file = 'SItRDmodel_ONSparams_noise_NB_NO-R_travel_TRUE_step_TRUE.npy'
    else:
        folder_path =  os.path.join(MODULE_DIR, 'out_SIRvsAGE_SuplementaryFig')
        full_fit_data_file = 'SynteticSItD_default_params_travel_TRUE_step_TRUE_' + str(SyntDataNum_file) + '.npy'
    data_filename = full_fit_data_file

    # Load data
    data = np.load(os.path.join(folder_path, data_filename ))
    data_S = data[:,0]
    data_Itot = data[:,1]
    data_R = data[:,2]
    data_Dreal = data[:,3]
    data_D = data[:,4] # noise data
    data_I = data[:,5:].T # transpose to get exactly the same shape as other code

    if len(data_R) < maxtime_fit:
        p.maxtime = len(data_R) -1
        maxtime_fit = len(data_R) -1
    else:
        p.maxtime = maxtime_fit

    # cut the data to the maxtime length:
    data_D = data_D[:p.maxtime+1]
    data_Dreal = data_Dreal[:p.maxtime+1]
    data_S_long = data_S
    data_S = data_S[:p.maxtime+1]
    data_Itot = data_Itot[:p.maxtime+1]
    data_R = data_R[:p.maxtime+1]
    data_I = data_I[:,:p.maxtime+1]

    # to get the same data and fit lengths as in Data_loader
    p.maxtime = p.maxtime + p.numeric_max_age + p.extra_days_to_simulate
    p.day_1st_death_after_150220 = 22

    if FitFull:
        filename = get_file_name_suffix(p, 'SimSItD-' + str(SyntDataNum_file) + rho_label, Noise_label + \
             'model-' + ModelName + '_full-fit-' + str(FitFull), parameters_to_optimise)
    else:
        filename = get_file_name_suffix_anymodel(p, 'SimSItD-' + str(SyntDataNum_file), rho_label, Noise_label, ModelName , parameters_to_optimise)
    filename = filename + fixed_params_tag

    # filename = filename + '-PopulationMCMC' #HaarioBardenetACMC EmceeHammerMCMC PopulationMCMC RaoBlackwellACMC
    folder = args.out_mcmc
    saveas = os.path.join(MODULE_DIR, folder, filename)
    chains = pints.io.load_samples(saveas + '-chain.csv', nchains)
    logpdfs = pints.io.load_samples(saveas + '-logpdf.csv', nchains) # file contains log-posterior, log-likelihood, log-prior
    chains = np.array(chains)
    logpdfs = np.array(logpdfs)

    model_func = MODEL_FUNCTIONS[ModelName]
    LL = noise_model(p, data_D[p.day_1st_death_after_150220:], parameters_to_optimise, model_func=model_func)

    # Discard burn in
    burn_in = args.burn_in
    chains = chains[:, burn_in:, :]
    logpdfs = logpdfs[:, burn_in:, :-1]
    print('Burn-in period = ' + str(args.burn_in) + ' iterations')

    print('Plots for log-likelihoods... ')
    pints.plot.trace(logpdfs, parameter_names=['log-posterior', 'log-L'])
    if not args.show_plots:
        plt.savefig(saveas + '_TraceLikelihoods.png')

    print('Plots for the parameter estimates...')
    pints.plot.trace(chains, parameter_names=parameter_names)
    if not args.show_plots:
        plt.savefig(saveas + '_chains.png')

    # Apply thinning
    chains = chains[:, ::10, :]
    logpdfs = logpdfs[:, ::10, :]

     # # Look at histograms
    pints.plot.histogram(chains, kde=True, n_percentiles=99, parameter_names=parameter_names)
    if not args.show_plots:
        plt.savefig(saveas + '_histograms.png')

    # Specific to the selected chain:
    pints.plot.pairwise(chains[args.chain - 1], kde=True, n_percentiles=99, parameter_names=parameter_names)
    if not args.show_plots:
        plt.savefig(saveas + '_pairwise_posteriors_chain' + str(args.chain) + '.png')

    print('Chains summary:')
    mcmc_summary = pints.MCMCSummary(chains, parameter_names=parameter_names).summary()
    print(*mcmc_summary, sep='\n')
    mcmc_logL_summary = pints.MCMCSummary(logpdfs, parameter_names=['log-posterior', 'log-L']).summary()
    print(*mcmc_logL_summary, sep='\n')

    # Show graphs
    if args.show_plots:
        plt.show()