from typing import TypedDict
from logging import Logger

import polars as pl


class HistoryPosition(TypedDict):
    balance: float
    position: pl.DataFrame


class Portfolio:
    def __init__(self, code_categories: pl.Enum,  logger: Logger, balance: float = 0, commission: float = 0):
        self.position: pl.DataFrame = pl.DataFrame(
            data={
                'code': pl.Series(dtype=code_categories), # 标的代码
                'name': pl.Series(dtype=pl.String), # 标的名称
                'volume': pl.Series(dtype=pl.Float64), # 持仓量
                'cost_price': pl.Series(dtype=pl.Float64), # 成本价
                'last_price': pl.Series(dtype=pl.Float64), # 最新价
                'market_value': pl.Series(dtype=pl.Float64), # 市值
                'suspension': pl.Series(dtype=pl.Boolean), # 是否停牌
            }
        )
        self.balance: float = balance # 现金
        self.commission: float = commission # 手续费
        self.history_position: dict[str, HistoryPosition] = {} # 历史持仓
        self.history_dividend: dict[str, pl.DataFrame] = {} # 历史分红
        self.logger = logger
        
    def __repr__(self):
        return (
            f"Portfolio(balance={self.balance}, "
            f"equity={self.equity}, "
            f"available_equity={self.available_equity}, "
            f"unavailable_equity={self.unavailable_equity})\n"
            f"{self.position}"
        )

    @property
    def equity(self) -> float:
        """
        权益
        """
        market_value = (
            self.position
                .select(
                    pl.col('market_value').sum()
                )
                .item()
        )
        return round(self.balance + market_value, 2)
    
    @property
    def available_equity(self) -> float:
        """
        可用权益
        """
        available_market_value = (
            self.position
                .filter(
                    ~pl.col('suspension')
                )
                .select(
                    pl.col('market_value').sum()
                )
                .item()
        )
        return round(self.balance + available_market_value, 2)
    
    @property
    def unavailable_equity(self) -> float:
        """
        不可用权益
        """
        return (
            self.position
                .filter(
                    pl.col('suspension')
                )
                .select(
                    pl.col('market_value').sum().round(2)
                )
                .item()
        )
        

    def fill_target(self, target_lf: pl.LazyFrame, stamp_tax: float = 0):
        """
        根据目标持仓更新持仓
        :param target_lf: 目标持仓表, 列包含code, name, last_price, target_volume
        :type target_lf: pl.LazyFrame
        :param stamp_tax: 印花税
        :type stamp_tax: float
        """
# target_lf example:
# ┌───────────┬──────────┬────────────┬───────────────┐
# │ code      ┆ name     ┆ last_price ┆ target_volume │
# │ ---       ┆ ---      ┆ ---        ┆ ---           │
# │ enum      ┆ str      ┆ f64        ┆ f64           │
# ╞═══════════╪══════════╪════════════╪═══════════════╡
# │ 128126.SZ ┆ 蓝盾转债  ┆ 191.98     ┆ 440.0         │
# └───────────┴──────────┴────────────┴───────────────┘
        # self.logger.debug(f"目标篮子: \n{target_lf.collect()}")

# _df example:
# ┌────────┬────────┬──────────┬────────┬────────┬────────┬────────┬────────┬────────┬────────┬────────┬────────┐
# │ suspen ┆ code   ┆ name     ┆ volume ┆ cost_p ┆ last_p ┆ market ┆ target ┆ target ┆ fill_v ┆ fill_a ┆ commis │
# │ sion   ┆ ---    ┆ ---      ┆ ---    ┆ rice   ┆ rice   ┆ _value ┆ _volum ┆ _cost_ ┆ olume  ┆ mount  ┆ sion   │
# │ ---    ┆ enum   ┆ str      ┆ f64    ┆ ---    ┆ ---    ┆ ---    ┆ e      ┆ price  ┆ ---    ┆ ---    ┆ ---    │
# │ bool   ┆        ┆          ┆        ┆ f64    ┆ f64    ┆ f64    ┆ ---    ┆ ---    ┆ f64    ┆ f64    ┆ f64    │
# │        ┆        ┆          ┆        ┆        ┆        ┆        ┆ f64    ┆ f64    ┆        ┆        ┆        │
# ╞════════╪════════╪══════════╪════════╪════════╪════════╪════════╪════════╪════════╪════════╪════════╪════════╡
# │ false  ┆ 128126 ┆ 蓝盾转债  ┆ 1560.0 ┆ 181.01 ┆ 191.98 ┆ 299488 ┆ 1680.0 ┆ 181.80 ┆ 120.0  ┆ 23037. ┆ 2.3    │
# │        ┆ .SZ    ┆          ┆        ┆ 8103   ┆        ┆ .8     ┆        ┆ 2464   ┆        ┆ 6      ┆        │
# │ …      ┆ …      ┆ …        ┆ …      ┆ …      ┆ …      ┆ …      ┆ …      ┆ …      ┆ …      ┆ …      ┆ …      │
# └────────┴────────┴──────────┴────────┴────────┴────────┴────────┴────────┴────────┴────────┴────────┴────────┘
        position_lf = self.position.lazy()

        _suspension_expr = pl.col('suspension').fill_null(False) # 停牌标志 (虽然target_lf中没有停牌标的，但可以视为非停牌)
        _volume_expr = pl.col('volume').fill_null(0) # 持仓量 (在target_lf中，且不在postion中的标的，持仓量记录为0)
        _last_price_expr = pl.coalesce(pl.col('last_price_right'), pl.col('last_price')) # 最新价 (如果target_lf中没有最新价，则使用position表中的最新价)
        _name_expr = pl.coalesce(pl.col('name_right'), pl.col('name')) # 标的名称 (优先使用target_lf中的name，若无则使用position中的name)
        _market_value_expr = (_last_price_expr * _volume_expr).round(2) # 市值
        _target_volume_expr = pl.when(_suspension_expr).then(_volume_expr).otherwise(pl.col('target_volume').fill_null(0)) # 目标量 (停牌标的目标量设为现量，非停牌标的目标量设为篮子目标量)
        _fill_volume_expr = _target_volume_expr - _volume_expr # 调仓量
        _fill_amount_expr = _last_price_expr * _fill_volume_expr # 调仓金额
        _commission_expr = (_fill_amount_expr.abs() * self.commission).round(2) # 手续费
        # 目标成本价
        _target_cost_price_expr = (
            pl.when(pl.col('cost_price').is_null())
            .then(
                # 首次买入，没有历史成本
                (_fill_amount_expr + _commission_expr) / _fill_volume_expr
            ).otherwise(
                pl.when(_target_volume_expr == 0)
                .then(pl.lit(None)) # 清仓卖出
                .otherwise(
                    pl.when(_fill_volume_expr > 0)
                    .then(
                        # 非首次买入
                        (pl.col('cost_price') * _volume_expr + _fill_amount_expr + _commission_expr)
                        / _target_volume_expr 
                    )
                    .otherwise(
                        # 非清仓卖出
                        (pl.col('cost_price') * _volume_expr + _fill_amount_expr + _commission_expr + (_fill_amount_expr * stamp_tax).round(2))
                        / _target_volume_expr 
                    )
                )
            )
        )

        # 目标篮子表合并持仓表
        _df = position_lf.join(
            target_lf,
            on='code',
            how='full',
            coalesce=True,
            maintain_order='right' # 保持目标篮子的rank排序（有性能损耗）
        ).with_columns(
            _suspension_expr.alias('suspension'),
            _volume_expr.alias('volume'),
            _target_cost_price_expr.alias('target_cost_price'),
            _last_price_expr.alias('last_price'),
            _name_expr.alias('name'),
            _market_value_expr.alias('market_value'),
            _target_volume_expr.alias('target_volume'),
            _fill_volume_expr.alias('fill_volume'),
            _fill_amount_expr.alias('fill_amount'),
            _commission_expr.alias('commission'),
        ).select(
            pl.exclude('last_price_right', 'name_right') # 排除由于join后产生的列
        ).collect()
        # self.logger.debug(f"调仓预览: \n{_df}")

        # 更新balance
        buy_amount = _df.select(
            pl.col('fill_amount').filter(pl.col('fill_amount') > 0).sum().round(2)
        ).item()
        # 卖出金额
        sell_amount = _df.select(
            pl.col('fill_amount').filter(pl.col('fill_amount') < 0).abs().sum().round(2)
        ).item()
        # 买入手续费
        buy_commission = _df.select(
            pl.col('commission').filter(pl.col('fill_amount') > 0).sum().round(2)
        ).item()
        # 卖出手续费
        sell_commission = _df.select(
            pl.col('commission').filter(pl.col('fill_amount') < 0).sum().round(2)
        ).item() + round(sell_amount * stamp_tax, 2)
        self.balance = round(self.balance + sell_amount - buy_amount - buy_commission - sell_commission, 2)
# position_df example:
# ┌────────────┬───────────┬──────────┬────────┬────────────┬────────────┬──────────────┐
# │ suspension ┆ code      ┆ name     ┆ volume ┆ cost_price ┆ last_price ┆ market_value │
# │ ---        ┆ ---       ┆ ---      ┆ ---    ┆ ---        ┆ ---        ┆ ---          │
# │ bool       ┆ enum      ┆ str      ┆ f64    ┆ f64        ┆ f64        ┆ f64          │
# ╞════════════╪═══════════╪══════════╪════════╪════════════╪════════════╪══════════════╡
# │ false      ┆ 128126.SZ ┆ 蓝盾转债  ┆ 440.0  ┆ 179.525909 ┆ 191.98     ┆ 84471.2      │
# └────────────┴───────────┴──────────┴────────┴────────────┴────────────┴──────────────┘
        # 更新持仓
        position_df = _df.select(
            pl.col('suspension'),
            pl.col('code'),
            pl.col('name'),
            pl.col('target_volume').alias('volume'), # 更新持仓量
            pl.col('target_cost_price').alias('cost_price'), # 更新成本价
            pl.col('last_price'),
            (pl.col('target_volume') * pl.col('last_price')).alias('market_value') # 更新市值
        ).filter(
            pl.col('volume') != 0
        )
        # self.logger.debug(f"最新持仓持仓: \n{position_df}")
        self.position = position_df
        return buy_amount, sell_amount, buy_commission, sell_commission
    
    def update_history_position(self, trade_date: str):
        self.history_position[trade_date] = {
            'balance': self.balance,
            'position': self.position.clone()
        }
