"""
WorldQuant Brain API Client Library

A Python client for interacting with the WorldQuant Brain platform API.
Supports authentication, alpha simulation, and result retrieval.
"""

import json
import getpass
import requests
from time import sleep
from os.path import expanduser
from typing import Optional


BASE_URL = "https://api.worldquantbrain.com"


class BrainClient:
    """Client for the WorldQuant Brain API."""

    def __init__(
        self,
        email: Optional[str] = None,
        password: Optional[str] = None,
        credentials_file: Optional[str] = None,
    ):
        """
        Initialize the client. Credentials priority:
          1. email + password passed directly
          2. credentials_file path (JSON: ["email", "password"])
          3. ~/.brain_credentials (if the file exists)
          4. Interactive prompt (asks at runtime)

        Args:
            email: Brain platform email.
            password: Brain platform password.
            credentials_file: Path to a JSON credentials file.
        """
        self._session = requests.Session()
        self._authenticated = False

        if email and password:
            self._session.auth = (email, password)
            return

        if credentials_file:
            with open(expanduser(credentials_file), "r") as f:
                self._session.auth = tuple(json.load(f))
            return

        default_file = expanduser("~/.brain_credentials")
        try:
            with open(default_file, "r") as f:
                self._session.auth = tuple(json.load(f))
        except FileNotFoundError:
            # Fall back to interactive prompt
            email = input("WorldQuant Brain email: ").strip()
            password = getpass.getpass("WorldQuant Brain password: ")
            self._session.auth = (email, password)

    @classmethod
    def login(cls, email: Optional[str] = None, password: Optional[str] = None) -> "BrainClient":
        """
        Create a BrainClient and authenticate in one step.

        If email/password are not provided, prompts interactively.

        Args:
            email: Brain platform email (optional).
            password: Brain platform password (optional).

        Returns:
            An authenticated BrainClient instance.

        Example:
            client = BrainClient.login()                          # interactive prompt
            client = BrainClient.login("me@email.com", "pass")   # direct
        """
        client = cls(email=email, password=password)
        client.authenticate()
        return client

    # -------------------------------------------------------------------------
    # Authentication
    # -------------------------------------------------------------------------

    def authenticate(self) -> dict:
        """
        Sign in and obtain a JWT token.

        Returns:
            Response JSON from the /authentication endpoint.

        Raises:
            requests.HTTPError: If authentication fails.
        """
        response = self._session.post(f"{BASE_URL}/authentication")
        response.raise_for_status()
        self._authenticated = True
        return response.json()

    # -------------------------------------------------------------------------
    # Simulations
    # -------------------------------------------------------------------------

    def simulate(self, expression: str, settings: dict = None, alpha_type: str = "REGULAR", regular: str = "close") -> "SimulationResult":
        """
        Submit an alpha expression for simulation.

        Args:
            expression: The alpha expression string.
            settings: Simulation settings dict. Defaults to standard equity settings.
            alpha_type: Type of simulation, e.g. "REGULAR".
            regular: Price field, e.g. "close".

        Returns:
            A SimulationResult object to poll for completion.

        Raises:
            requests.HTTPError: If the submission fails.
        """
        default_settings = {
            "instrumentType": "EQUITY",
            "region": "USA",
            "universe": "TOP3000",
            "delay": 1,
            "decay": 15,
            "neutralization": "SUBINDUSTRY",
            "truncation": 0.08,
            "maxTrade": "ON",
            "pasteurization": "ON",
            "testPeriod": "P1Y6M",
            "unitHandling": "VERIFY",
            "nanHandling": "OFF",
            "language": "FASTEXPR",
            "visualization": False,
        }
        if settings:
            default_settings.update(settings)

        payload = {
            "type": alpha_type,
            "settings": default_settings,
            "regular": expression if regular == "close" else regular,
        }

        # If expression is provided separately from regular field
        if regular == "close":
            payload["regular"] = expression

        response = self._session.post(f"{BASE_URL}/simulations", json=payload)
        response.raise_for_status()

        progress_url = response.headers.get("Location")
        return SimulationResult(self._session, progress_url)

    def simulate_parallel(
        self,
        expressions: list,
        settings: dict = None,
        workers: int = 4,
        verbose: bool = True,
    ) -> list:
        """
        Submit multiple alpha expressions for simulation in parallel.

        Args:
            expressions: List of expression strings, or list of dicts with keys
                         'expression' (required) and 'settings' (optional override).
            settings: Shared simulation settings applied to all expressions (unless
                      overridden per-expression via dict input).
            workers: Number of parallel threads for submitting and polling.
            verbose: If True, print progress per expression.

        Returns:
            List of SimulationResult objects (same order as input), all completed.

        Example:
            results = client.simulate_parallel(
                ["ts_mean(returns, 5)", "ts_sum(volume, 10)"],
                settings={"region": "USA", "universe": "TOP3000"},
                workers=2,
            )
            for r in results:
                alpha = r.get_alpha()
                print(r.alpha_id, alpha["is"]["sharpe"])
        """
        from concurrent.futures import ThreadPoolExecutor, as_completed

        def _run(item):
            if isinstance(item, dict):
                expr = item["expression"]
                merged = {**(settings or {}), **item.get("settings", {})}
                sim = self.simulate(expr, merged or None)
            else:
                sim = self.simulate(item, settings)
            sim.wait(verbose=verbose)
            return sim

        results_map = {}
        with ThreadPoolExecutor(max_workers=workers) as pool:
            futures = {pool.submit(_run, item): idx for idx, item in enumerate(expressions)}
            for future in as_completed(futures):
                idx = futures[future]
                results_map[idx] = future.result()

        return [results_map[i] for i in range(len(expressions))]

    # -------------------------------------------------------------------------
    # Alphas
    # -------------------------------------------------------------------------

    def get_alpha(self, alpha_id: str) -> dict:
        """
        Retrieve details of a completed alpha by ID.

        Args:
            alpha_id: The alpha ID string.

        Returns:
            Alpha details as a dict.
        """
        response = self._session.get(f"{BASE_URL}/alphas/{alpha_id}")
        response.raise_for_status()
        return response.json()

    def get_pnl(self, alpha_id: str, poll_interval: float = 5.0) -> dict:
        """
        Retrieve PnL record set for an alpha, polling until ready.

        Args:
            alpha_id: The alpha ID string.
            poll_interval: Fallback seconds to wait between polls if no Retry-After header.

        Returns:
            PnL record set as a dict.
        """
        return self._poll_recordset(alpha_id, "pnl", poll_interval)

    def get_recordset(self, alpha_id: str, record_set_name: str, poll_interval: float = 5.0) -> dict:
        """
        Retrieve any named record set for an alpha, polling until ready.

        Args:
            alpha_id: The alpha ID string.
            record_set_name: Name of the record set (e.g. "pnl", "sharpe").
            poll_interval: Fallback seconds between polls if no Retry-After header.

        Returns:
            Record set data as a dict.
        """
        return self._poll_recordset(alpha_id, record_set_name, poll_interval)

    def get_yearly(self, alpha_id: str, poll_interval: float = 5.0) -> list:
        """
        Retrieve PnL grouped by year for an alpha.

        Args:
            alpha_id: The alpha ID string.
            poll_interval: Fallback seconds between polls if no Retry-After header.

        Returns:
            List of dicts, one per year, with keys: year, pnl, returns, sharpe, turnover, drawdown (when available).
        """
        from collections import defaultdict
        import math

        pnl_data = self.get_pnl(alpha_id, poll_interval)
        # records = [[date, pnl], ...] — schema: properties[0]=date, properties[1]=pnl
        records = pnl_data.get("records", [])

        grouped = defaultdict(list)
        for r in records:
            year = str(r[0])[:4]  # r[0] = date string "YYYY-MM-DD"
            if year.isdigit():
                grouped[year].append(float(r[1]))  # r[1] = daily pnl

        result = []
        for year in sorted(grouped.keys()):
            daily_pnls = grouped[year]
            n = len(daily_pnls)
            total_pnl = sum(daily_pnls)

            # Sharpe = mean(daily_pnl) / std(daily_pnl) * sqrt(252)
            if n > 1:
                mean = total_pnl / n
                variance = sum((x - mean) ** 2 for x in daily_pnls) / (n - 1)
                std = math.sqrt(variance) if variance > 0 else 0
                sharpe = round(mean / std * math.sqrt(252) / 100, 4) if std > 0 else None
            else:
                sharpe = None

            result.append({
                "year": int(year),
                "pnl": round(total_pnl, 2),
                "sharpe": sharpe,
                "days": n,
            })

        return result

    def _poll_recordset(self, alpha_id: str, record_set_name: str, poll_interval: float) -> dict:
        url = f"{BASE_URL}/alphas/{alpha_id}/recordsets/{record_set_name}"
        while True:
            response = self._session.get(url)
            retry_after = float(response.headers.get("Retry-After", 0))
            if retry_after == 0:
                response.raise_for_status()
                return response.json()
            #print(f"Sleeping for {retry_after} seconds...")
            sleep(retry_after)


class SimulationResult:
    """Represents a pending or completed alpha simulation."""

    def __init__(self, session: requests.Session, progress_url: str):
        self._session = session
        self.progress_url = progress_url
        self.alpha_id: str = None
        self._result: dict = None

    def wait(self, verbose: bool = True) -> dict:
        """
        Poll until the simulation completes.

        Args:
            verbose: If True, print polling status messages.

        Returns:
            The completed simulation result JSON containing the alpha id.
        """
        while True:
            response = self._session.get(self.progress_url)
            retry_after = float(response.headers.get("Retry-After", 0))
            if retry_after == 0:
                response.raise_for_status()
                self._result = response.json()
                self.alpha_id = self._result.get("alpha")
                if verbose:
                    print(f"Alpha simulation complete. Alpha ID: {self.alpha_id}")
                return self._result
            if verbose:
                print(f"Simulating... sleeping for {retry_after} seconds")
            sleep(retry_after)

    def get_alpha(self) -> dict:
        """
        Fetch full alpha details after simulation completes.

        Returns:
            Alpha details dict.

        Raises:
            RuntimeError: If wait() has not been called yet.
        """
        if not self.alpha_id:
            raise RuntimeError("Simulation not complete. Call wait() first.")
        response = self._session.get(f"{BASE_URL}/alphas/{self.alpha_id}")
        response.raise_for_status()
        return response.json()

    def get_pnl(self, poll_interval: float = 5.0) -> dict:
        """
        Retrieve PnL record set after simulation completes.

        Args:
            poll_interval: Fallback seconds between polls if no Retry-After header.

        Returns:
            PnL record set dict.

        Raises:
            RuntimeError: If wait() has not been called yet.
        """
        if not self.alpha_id:
            raise RuntimeError("Simulation not complete. Call wait() first.")
        client = BrainClient.__new__(BrainClient)
        client._session = self._session
        client._authenticated = True
        return client.get_pnl(self.alpha_id, poll_interval)

    def get_yearly(self, poll_interval: float = 5.0) -> list:
        """
        Retrieve PnL grouped by year after simulation completes.

        Returns:
            List of dicts, one per year.

        Raises:
            RuntimeError: If wait() has not been called yet.
        """
        if not self.alpha_id:
            raise RuntimeError("Simulation not complete. Call wait() first.")
        client = BrainClient.__new__(BrainClient)
        client._session = self._session
        client._authenticated = True
        return client.get_yearly(self.alpha_id, poll_interval)
