"""The lab's fixed adapter. Visitor input is data, never Python source."""
import json
import pandas as pd
from quantcore.portfolio import portfolio_returns, portfolio_volatility, correlation_matrix
from quantcore.returns import wealth_index
from quantcore.risk import max_drawdown


def compare(parameters, datasets):
    if not isinstance(parameters, dict) or set(parameters) != {"dataset", "left", "right"}:
        raise ValueError("Choose a dataset and two allocations.")
    if parameters["dataset"] not in ("diverse", "correlated"):
        raise ValueError("Choose one of the two sample datasets.")
    for key in ("left", "right"):
        n = parameters[key]
        if isinstance(n, bool) or not isinstance(n, int) or not 0 <= n <= 100:
            raise ValueError("Asset A must be a whole percentage from 0 to 100.")
    frame = pd.DataFrame(datasets[parameters["dataset"]], columns=["A", "B"])

    def calculate(a):
        weights = pd.Series({"A": a / 100, "B": 1 - a / 100})
        returns = portfolio_returns(frame, weights)
        wealth = wealth_index(returns, initial_value=100)
        # Include day zero when finding drawdown so an initial loss is counted.
        path = pd.concat([pd.Series([100.0]), wealth], ignore_index=True)
        return {"aWeight": a, "totalReturn": float(path.iloc[-1] / 100 - 1),
                "volatility": float(portfolio_volatility(frame, weights)),
                "maxDrawdown": float(max_drawdown(path)), "path": path.tolist()}

    return {"dataset": parameters["dataset"], "left": calculate(parameters["left"]),
            "right": calculate(parameters["right"]),
            "correlation": float(correlation_matrix(frame).loc["A", "B"]),
            "observations": len(frame)}


def compare_json(parameters_json, datasets_json):
    return json.dumps(compare(json.loads(parameters_json), json.loads(datasets_json)), allow_nan=False)