use std::collections::BTreeSet;

use chrono::NaiveDate;
use serde::Serialize;

use crate::data::{BenchmarkSnapshot, DataSet, EligibleUniverseSnapshot};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BandRegime {
    Bullish,
    Neutral,
    Defensive,
}

#[derive(Debug, Clone, Serialize)]
pub struct UniverseCandidate {
    pub symbol: String,
    pub market_cap_bn: f64,
    pub free_float_cap_bn: f64,
    pub band_low: f64,
    pub band_high: f64,
}

#[derive(Debug, Clone, Serialize)]
pub struct SelectionDiagnostics {
    pub decision_date: NaiveDate,
    pub reference_level: f64,
    pub band_low: f64,
    pub band_high: f64,
    pub factor_total: usize,
    pub market_cap_missing_count: usize,
    pub candidate_missing_count: usize,
    pub market_missing_count: usize,
    pub not_eligible_count: usize,
    pub paused_count: usize,
    pub out_of_band_count: usize,
    pub selected_before_limit: usize,
    pub selected_after_limit: usize,
    pub missing_market_cap_symbols: Vec<String>,
    pub selected_symbols: Vec<String>,
    pub rejection_examples: Vec<String>,
}

pub struct SelectionContext<'a> {
    pub decision_date: NaiveDate,
    pub benchmark: &'a BenchmarkSnapshot,
    pub reference_level: f64,
    pub data: &'a DataSet,
    pub dynamic_universe: Option<&'a BTreeSet<String>>,
}

impl SelectionContext<'_> {
    fn eligible_universe(&self) -> Vec<EligibleUniverseSnapshot> {
        let eligible = self.data.eligible_universe_on(self.decision_date);
        match self.dynamic_universe {
            Some(symbols) if !symbols.is_empty() => eligible
                .iter()
                .filter(|row| symbols.contains(&row.symbol))
                .cloned()
                .collect(),
            _ => eligible.to_vec(),
        }
    }
}

pub trait UniverseSelector {
    fn select(&self, ctx: &SelectionContext<'_>) -> Vec<UniverseCandidate>;
    fn select_with_diagnostics(
        &self,
        ctx: &SelectionContext<'_>,
    ) -> (Vec<UniverseCandidate>, SelectionDiagnostics);
}

#[derive(Debug, Clone)]
pub struct DynamicMarketCapBandSelector {
    pub base_index_level: f64,
    pub base_cap_floor: f64,
    pub cap_span: f64,
    pub xs: f64,
    pub top_n: usize,
    pub padding_ratio: f64,
    pub min_padding: f64,
    pub max_padding: f64,
}

impl DynamicMarketCapBandSelector {
    pub fn new(
        base_index_level: f64,
        base_cap_floor: f64,
        cap_span: f64,
        xs: f64,
        top_n: usize,
        padding_ratio: f64,
        min_padding: f64,
        max_padding: f64,
    ) -> Self {
        Self {
            base_index_level,
            base_cap_floor,
            cap_span,
            xs,
            top_n,
            padding_ratio,
            min_padding,
            max_padding,
        }
    }

    pub fn demo(top_n: usize) -> Self {
        Self::new(2000.0, 7.0, 10.0, 4.0 / 500.0, top_n, 0.5, 8.0, 20.0)
    }

    pub fn regime(&self, benchmark_level: f64) -> BandRegime {
        if benchmark_level >= self.base_index_level + 400.0 {
            BandRegime::Bullish
        } else if benchmark_level >= self.base_index_level {
            BandRegime::Neutral
        } else {
            BandRegime::Defensive
        }
    }

    pub fn band_for_level(&self, benchmark_level: f64) -> (f64, f64) {
        let start = ((benchmark_level - self.base_index_level) * self.xs) + self.base_cap_floor;
        let low = start.round();
        let high = low + self.cap_span;

        // Apply padding to expand the range
        let span = high - low;
        let padding = (span * self.padding_ratio)
            .max(self.min_padding)
            .min(self.max_padding);

        let lower_bound = (low - padding).max(0.0);
        let upper_bound = high + padding;

        (lower_bound, upper_bound)
    }
}

impl UniverseSelector for DynamicMarketCapBandSelector {
    fn select(&self, ctx: &SelectionContext<'_>) -> Vec<UniverseCandidate> {
        self.select_with_diagnostics(ctx).0
    }

    fn select_with_diagnostics(
        &self,
        ctx: &SelectionContext<'_>,
    ) -> (Vec<UniverseCandidate>, SelectionDiagnostics) {
        let _regime = self.regime(ctx.reference_level);
        let (min_cap, max_cap) = self.band_for_level(ctx.reference_level);
        let mut diagnostics = SelectionDiagnostics {
            decision_date: ctx.decision_date,
            reference_level: ctx.reference_level,
            band_low: min_cap,
            band_high: max_cap,
            factor_total: 0,
            market_cap_missing_count: 0,
            candidate_missing_count: 0,
            market_missing_count: 0,
            not_eligible_count: 0,
            paused_count: 0,
            out_of_band_count: 0,
            selected_before_limit: 0,
            selected_after_limit: 0,
            missing_market_cap_symbols: Vec::new(),
            selected_symbols: Vec::new(),
            rejection_examples: Vec::new(),
        };

        diagnostics.factor_total = ctx.data.factor_snapshots_on(ctx.decision_date).len();
        let eligible = ctx.eligible_universe();
        diagnostics.market_cap_missing_count =
            diagnostics.factor_total.saturating_sub(eligible.len());
        let start_idx = lower_bound_by_market_cap(&eligible, min_cap);
        let mut selected = Vec::new();

        for factor in eligible.iter().skip(start_idx) {
            if factor.market_cap_bn > max_cap {
                break;
            }
            selected.push(to_universe_candidate(factor, min_cap, max_cap));
        }

        diagnostics.out_of_band_count = eligible.len().saturating_sub(selected.len());

        diagnostics.selected_before_limit = selected.len();
        if selected.len() > self.top_n {
            selected.truncate(self.top_n);
        }
        diagnostics.selected_symbols = selected.iter().map(|item| item.symbol.clone()).collect();
        diagnostics.selected_after_limit = diagnostics.selected_symbols.len();
        (selected, diagnostics)
    }
}

fn lower_bound_by_market_cap(rows: &[EligibleUniverseSnapshot], target: f64) -> usize {
    let mut left = 0usize;
    let mut right = rows.len();
    while left < right {
        let mid = left + (right - left) / 2;
        if rows[mid].market_cap_bn < target {
            left = mid + 1;
        } else {
            right = mid;
        }
    }
    left
}

fn to_universe_candidate(
    factor: &EligibleUniverseSnapshot,
    band_low: f64,
    band_high: f64,
) -> UniverseCandidate {
    UniverseCandidate {
        symbol: factor.symbol.clone(),
        market_cap_bn: factor.market_cap_bn,
        free_float_cap_bn: factor.free_float_cap_bn,
        band_low,
        band_high,
    }
}