from __future__ import annotations

import enum
import logging
from typing import NamedTuple, cast

import numpy as np
import pandas as pd

from signalab.units import Lot, Units

logger = logging.getLogger(__name__)

# Column indices in the positions matrix
_COL_IS_BUY = 0
_COL_OPEN_ASK = 1
_COL_OPEN_BID = 2
_COL_SIZE = 3
_COL_VALUE = 4
_COL_MARGIN = 5
_N_COLS = 6

NUMPY_THRESHOLD = 20


class Action(enum.Enum):
    BUY = "BUY"
    SELL = "SELL"


class SimpleMarketData(NamedTuple):
    Index: pd.Timestamp
    bid: float
    ask: float


class InssufficientMarginError(RuntimeError):
    pass


class Account:
    def __init__(
        self,
        initial_balance: float,
        leverage: float = 1.0,
        liquidation_level: float = 0.5,
        label: str = "",
    ):
        self.initial_balance = initial_balance
        self.balance = initial_balance
        self.leverage = leverage
        self.liquidation_level = liquidation_level
        self.label = label if label else f"Account_{id(self)}"

        self._pos: np.ndarray = np.empty((0, _N_COLS), dtype=np.float64)
        self._pos_ids: list[int] = []
        self._next_id: int = 0

        self._cached_used_margin: float = 0.0

        # Logging statistics
        self._realized_pl: float = 0.0
        self.opened_buy_count: int = 0
        self.opened_sell_count: int = 0
        self.opening_net_impact: float = 0.0
        self.maintain_buy_net_impact: float = 0.0
        self.maintain_sell_net_impact: float = 0.0

        # self.record: dict[str, np.ndarray] = {}
        self._record_timestamps: np.ndarray
        self._record_balance: np.ndarray
        self._record_equity: np.ndarray
        self._record_used_margin: np.ndarray
        self._record_free_margin: np.ndarray
        self._record_floating_pl: np.ndarray
        self._record_realized_pl: np.ndarray
        self._record_open_buy_count: np.ndarray
        self._record_open_sell_count: np.ndarray
        self._record_opening_cost: np.ndarray
        self._record_maintain_buy_cost: np.ndarray
        self._record_maintain_sell_cost: np.ndarray
        self._record_real_profit: np.ndarray

    def init_record(self, market_data: pd.DataFrame) -> None:
        length = len(market_data)
        self._record_timestamps = market_data.index.to_numpy()
        self._record_balance = np.zeros(length, dtype=np.float64)
        self._record_equity = np.zeros(length, dtype=np.float64)
        self._record_used_margin = np.zeros(length, dtype=np.float64)
        self._record_free_margin = np.zeros(length, dtype=np.float64)
        self._record_floating_pl = np.zeros(length, dtype=np.float64)
        self._record_realized_pl = np.zeros(length, dtype=np.float64)
        self._record_open_buy_count = np.zeros(length, dtype=np.float64)
        self._record_open_sell_count = np.zeros(length, dtype=np.float64)
        self._record_opening_cost = np.zeros(length, dtype=np.float64)
        self._record_maintain_buy_cost = np.zeros(length, dtype=np.float64)
        self._record_maintain_sell_cost = np.zeros(length, dtype=np.float64)
        self._record_real_profit = np.zeros(length, dtype=np.float64)

    def get_record_df(self) -> pd.DataFrame:
        record_df = pd.DataFrame(
            {
                "timestamp": self._record_timestamps,
                "balance": self._record_balance,
                "equity": self._record_equity,
                "used_margin": self._record_used_margin,
                "free_margin": self._record_free_margin,
                "floating_pl": self._record_floating_pl,
                "realized_pl": self._record_realized_pl,
                "open_buy_count": self._record_open_buy_count,
                "open_sell_count": self._record_open_sell_count,
                "opening_net_impact": self._record_opening_cost,
                "maintain_buy_net_impact": self._record_maintain_buy_cost,
                "maintain_sell_net_impact": self._record_maintain_sell_cost,
                "real_profit": self._record_real_profit,
            }
        ).set_index("timestamp")

        return record_df

    def floating_pl(self, market: SimpleMarketData) -> float:
        if len(self._pos_ids) == 0:
            return 0.0
        bid = market.bid
        ask = market.ask
        is_buy = self._pos[:, _COL_IS_BUY].astype(bool)
        open_asks = self._pos[:, _COL_OPEN_ASK]
        open_bids = self._pos[:, _COL_OPEN_BID]
        values = self._pos[:, _COL_VALUE]
        pls = np.where(is_buy, (bid - open_asks) * values, (open_bids - ask) * values)
        return float(pls.sum())

    def positions_unified_size(self) -> float:
        """Returns the total position size in lots."""
        if len(self._pos_ids) == 0:
            return 0.0
        sizes = self._pos[:, _COL_SIZE]
        return float(sizes.sum())

    def used_margin(self) -> float:
        return self._cached_used_margin

    def equity(self, market: SimpleMarketData) -> float:
        return self.balance + self.floating_pl(market)

    def free_margin(self, market: SimpleMarketData) -> float:
        return self.equity(market) - self._cached_used_margin

    def realized_pl(self) -> float:
        return self._realized_pl

    def open(
        self,
        action: Action,
        market_at_open: SimpleMarketData,
        size: Lot,
        leverage: float | None = None,
    ) -> int:
        lev = leverage or self.leverage
        value = size * Units.LOT_SIZE
        margin = value / lev

        if margin > self.free_margin(market_at_open):
            raise InssufficientMarginError(
                f"Cannot open position: required margin {margin:.2f} exceeds"
                f" free margin {self.free_margin(market_at_open):.2f}. You can"
                "check free margin with account.free_margin(market)."
            )

        is_buy = 1.0 if action == Action.BUY else 0.0
        row = np.array(
            [[is_buy, market_at_open.ask, market_at_open.bid, size, value, margin]]
        )
        self._pos = np.vstack([self._pos, row]) if len(self._pos_ids) > 0 else row
        pos_id = self._next_id
        self._next_id += 1
        self._pos_ids.append(pos_id)

        # Update cached margin
        self._cached_used_margin += margin

        # Apply commission
        commission = Units.open_commission(size)
        self.balance += commission
        self.opening_net_impact += commission

        if action == Action.BUY:
            self.opened_buy_count += 1
        else:
            self.opened_sell_count += 1

        return pos_id

    def close(self, pos_id: int, market_at_close: SimpleMarketData) -> bool:
        if pos_id not in self._pos_ids:
            return False

        idx = self._pos_ids.index(pos_id)
        row = self._pos[idx]

        is_buy = bool(row[_COL_IS_BUY])
        open_ask = row[_COL_OPEN_ASK]
        open_bid = row[_COL_OPEN_BID]
        value = row[_COL_VALUE]
        margin = row[_COL_MARGIN]

        if is_buy:
            pl = (market_at_close.bid - open_ask) * value
        else:
            pl = (open_bid - market_at_close.ask) * value

        self.balance += pl
        self._realized_pl += pl
        self._cached_used_margin -= margin

        # Remove from arrays
        self._pos = np.delete(self._pos, idx, axis=0)
        self._pos_ids.pop(idx)
        return True

    def close_all_positions(self, market: SimpleMarketData) -> None:
        # Vectorized P&L for all positions at once
        if len(self._pos_ids) == 0:
            return

        is_buy = self._pos[:, _COL_IS_BUY].astype(bool)
        open_asks = self._pos[:, _COL_OPEN_ASK]
        open_bids = self._pos[:, _COL_OPEN_BID]
        values = self._pos[:, _COL_VALUE]

        pls = np.where(
            is_buy,
            (market.bid - open_asks) * values,
            (open_bids - market.ask) * values,
        )

        total_pl = float(pls.sum())
        self.balance += total_pl
        self._realized_pl += total_pl
        self._cached_used_margin = 0.0
        self._pos = np.empty((0, _N_COLS), dtype=np.float64)
        self._pos_ids.clear()

    def check_liquidation(self, market: SimpleMarketData) -> None:
        if self._cached_used_margin == 0:
            return
        margin_level = self.equity(market) / self._cached_used_margin
        if margin_level < self.liquidation_level:
            self.close_all_positions(market)

    def day_change(self, delta: int) -> None:
        if len(self._pos_ids) == 0:
            return

        is_buy = self._pos[:, _COL_IS_BUY].astype(bool)
        sizes = self._pos[:, _COL_SIZE]

        buy_delta = Units.maintain_buy_commission(float((sizes[is_buy]).sum())) * delta
        sell_delta = (
            Units.maintain_sell_commission(float((sizes[~is_buy]).sum())) * delta
        )

        self.balance += buy_delta + sell_delta
        self.maintain_buy_net_impact += buy_delta
        self.maintain_sell_net_impact += sell_delta

    @property
    def open_position_count(self) -> int:
        return len(self._pos_ids)

    def position_pls(self, market: SimpleMarketData) -> dict[int, float]:
        """
        Returns a dict mapping pos_id -> current floating P&L.
        Computed in one vectorized pass — use this in strategies
        instead of iterating positions manually.

        Example:
            for pos_id, pl in self.account.position_pls(market).items():
                if pl > threshold:
                    self.account.close(pos_id, market)
        """
        if not self._pos_ids:
            return {}

        if len(self._pos_ids) >= NUMPY_THRESHOLD:
            is_buy = self._pos[:, _COL_IS_BUY].astype(bool)
            values = self._pos[:, _COL_VALUE]
            pls = np.where(
                is_buy,
                (market.bid - self._pos[:, _COL_OPEN_ASK]) * values,
                (self._pos[:, _COL_OPEN_BID] - market.ask) * values,
            )
            return dict(zip(self._pos_ids, pls.tolist(), strict=True))

        bid = market.bid
        ask = market.ask
        result = {}
        for i, pos_id in enumerate(self._pos_ids):
            row = self._pos[i]
            if row[_COL_IS_BUY]:
                result[pos_id] = (bid - row[_COL_OPEN_ASK]) * row[_COL_VALUE]
            else:
                result[pos_id] = (row[_COL_OPEN_BID] - ask) * row[_COL_VALUE]
        return result

    def position_size(self, pos_id: int) -> float:
        if pos_id not in self._pos_ids:
            raise ValueError(f"Position ID {pos_id} not found.")
        idx = self._pos_ids.index(pos_id)
        return cast(float, self._pos[idx, _COL_SIZE])

    def get_open_position_ids(self) -> list[int]:
        return list(self._pos_ids)

    def update_record(self, index: int, current_market: SimpleMarketData) -> None:
        try:
            self._record_balance[index] = self.balance
            self._record_equity[index] = self.equity(current_market)
            self._record_used_margin[index] = self._cached_used_margin
            self._record_free_margin[index] = self.free_margin(current_market)
            self._record_floating_pl[index] = self.floating_pl(current_market)
            self._record_realized_pl[index] = self._realized_pl
            self._record_open_buy_count[index] = float(self.opened_buy_count)
            self._record_open_sell_count[index] = float(self.opened_sell_count)
            self._record_opening_cost[index] = self.opening_net_impact
            self._record_maintain_buy_cost[index] = self.maintain_buy_net_impact
            self._record_maintain_sell_cost[index] = self.maintain_sell_net_impact
            self._record_real_profit[index] = (
                self._realized_pl
                + self.opening_net_impact
                + self.maintain_buy_net_impact
                + self.maintain_sell_net_impact
            )
        except AttributeError as err:
            raise RuntimeError(
                "Record arrays not initialized. Call init_record() before updating record."
            ) from err

    def copy_empty(self) -> Account:
        return Account(
            initial_balance=self.balance,
            leverage=self.leverage,
            liquidation_level=self.liquidation_level,
        )
