use crate::record::{
    EvaluationRecord, ProblemRecord, SolverRecord, TrialRecord, TrialRecordBuilder,
};
use crate::study::{Scheduling, StudyRecipe};
use crate::time::DateTime;
use chrono::Local;
use kurobako_core::hypervolume;
use kurobako_core::num::OrderedFloat;
use kurobako_core::problem::ProblemSpec;
use kurobako_core::solver::SolverSpec;
use kurobako_core::trial::{Params, TrialId, Values};
use kurobako_core::{Error, Result};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::cmp::Ordering;
use std::collections::{BTreeMap, BinaryHeap};
use std::fmt::Write as _;
use std::num::NonZeroUsize;
use std::time::Duration;

#[derive(Debug)]
pub struct StudyRecordBuilder {
    recipe: StudyRecipe,
    solver: SolverSpec,
    problem: ProblemSpec,
    start_time: DateTime,
    trials: BTreeMap<TrialId, TrialRecord>,
    pareto_frontier: BTreeMap<TrialId, (Params, Values)>,
}
impl StudyRecordBuilder {
    pub fn new(recipe: StudyRecipe, solver: SolverSpec, problem: ProblemSpec) -> Self {
        Self {
            recipe,
            solver,
            problem,
            start_time: Local::now(),
            trials: BTreeMap::new(),
            pareto_frontier: BTreeMap::new(),
        }
    }

    pub fn add_trial(&mut self, trial: TrialRecordBuilder) {
        let t = self.trials.entry(trial.id).or_insert_with(|| TrialRecord {
            thread_id: trial.thread_id,
            params: trial.params.clone(),
            evaluations: Vec::new(),
        });

        t.evaluations.push(EvaluationRecord {
            values: trial.values.clone(),
            start_step: trial.start_step,
            end_step: trial.end_step,
            ask_elapsed: trial.ask_elapsed,
            tell_elapsed: trial.tell_elapsed,
            evaluate_elapsed: trial.evaluate_elapsed,
        });

        if t.steps() == self.problem.steps.last() {
            let is_dominated = self
                .pareto_frontier
                .values()
                .any(|(_, vs)| vs.partial_cmp(&trial.values) == Some(Ordering::Less));
            if !is_dominated {
                let dominated = self
                    .pareto_frontier
                    .iter()
                    .filter(|(_, (_, vs))| trial.values.partial_cmp(vs) == Some(Ordering::Less))
                    .map(|(&id, _)| id)
                    .collect::<Vec<_>>();

                self.pareto_frontier
                    .insert(trial.id, (trial.params, trial.values));
                for id in dominated {
                    self.pareto_frontier.remove(&id);
                }
            }
        }
    }

    pub fn pareto_frontier(&self) -> impl '_ + Iterator<Item = (TrialId, &Params, &Values)> {
        self.pareto_frontier
            .iter()
            .map(|(&id, (params, values))| (id, params, values))
    }

    pub fn finish(self) -> StudyRecord {
        StudyRecord {
            start_time: self.start_time,
            end_time: Local::now(),
            budget: self.recipe.budget,
            seed: self.recipe.seed.unwrap_or_else(|| unreachable!()),
            concurrency: self.recipe.concurrency,
            scheduling: self.recipe.scheduling,
            solver: SolverRecord {
                recipe: self.recipe.solver,
                spec: self.solver,
            },
            problem: ProblemRecord {
                recipe: self.recipe.problem,
                spec: self.problem,
            },
            trials: self.trials.into_iter().map(|(_, v)| v).collect(),
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StudyRecord {
    pub start_time: DateTime,
    pub end_time: DateTime,
    pub seed: u64,
    pub budget: u64,
    pub concurrency: NonZeroUsize,
    pub scheduling: Scheduling,
    pub solver: SolverRecord,
    pub problem: ProblemRecord,
    pub trials: Vec<TrialRecord>,
}
impl StudyRecord {
    pub fn id(&self) -> Result<String> {
        let mut hasher = Sha256::new();
        hasher.update(&track!(
            serde_json::to_vec(&self.budget).map_err(Error::from)
        )?);
        hasher.update(&track!(
            serde_json::to_vec(&self.concurrency).map_err(Error::from)
        )?);
        hasher.update(&track!(
            serde_json::to_vec(&self.scheduling).map_err(Error::from)
        )?);
        hasher.update(&track!(
            serde_json::to_vec(&self.solver).map_err(Error::from)
        )?);
        hasher.update(&track!(
            serde_json::to_vec(&self.problem).map_err(Error::from)
        )?);

        let mut id = String::with_capacity(64);
        for b in hasher.finalize().as_slice() {
            track_write!(&mut id, "{:02x}", b)?;
        }
        Ok(id)
    }

    pub fn study_steps(&self) -> u64 {
        self.problem.spec.steps.last() * self.budget
    }

    pub fn best_values(&self) -> BTreeMap<u64, f64> {
        let mut best_values = BTreeMap::new();

        let problem_steps = self.problem.spec.steps.last();
        let mut trials = self
            .trials
            .iter()
            .filter_map(|t| {
                if let (Some(step), Some(value)) = (t.end_step(), t.value(problem_steps)) {
                    Some((step, value))
                } else {
                    None
                }
            })
            .collect::<Vec<_>>();
        trials.sort_by_key(|t| t.0);

        let mut min = std::f64::INFINITY;
        for (step, value) in trials {
            if value < min {
                min = value;
                best_values.insert(step, min);
            }
        }

        best_values
    }

    pub fn hypervolumes(&self) -> BTreeMap<u64, f64> {
        let mut hypervolumes = BTreeMap::new();

        let problem_steps = self.problem.spec.steps.last();
        let mut trials = self
            .trials
            .iter()
            .filter_map(|t| {
                if let (Some(step), Some(value)) = (t.end_step(), t.values(problem_steps)) {
                    Some((step, value))
                } else {
                    None
                }
            })
            .collect::<Vec<_>>();

        trials.sort_by_key(|t| t.0);

        let ref_pt = match &self.problem.spec.reference_point {
            Some(reference_point) => reference_point.to_vec(),
            None => vec![100.0; self.problem.spec.values_domain.len()],
        };

        let mut pts = Vec::new();
        for (step, values) in trials {
            pts.push(values.to_vec());
            let hv = hypervolume::compute(&pts, &ref_pt);
            hypervolumes.insert(step, hv);
        }

        hypervolumes
    }

    pub fn elapsed_times(&self, include_evaluate_time: bool) -> BTreeMap<u64, f64> {
        let mut times = BTreeMap::new();
        let mut elapsed = 0.0;
        for e in self.evaluations() {
            elapsed += e.ask_elapsed.get() + e.tell_elapsed.get();
            if include_evaluate_time {
                elapsed += e.evaluate_elapsed.get();
            }
            times.insert(e.end_step, elapsed);
        }
        times
    }

    pub fn evaluations(&self) -> impl '_ + Iterator<Item = &EvaluationRecord> {
        struct Entry<'a> {
            trial: &'a TrialRecord,
            index: usize,
        }
        impl<'a> Entry<'a> {
            fn step(&self) -> u64 {
                self.trial.evaluations[self.index].end_step
            }
        }
        impl<'a> PartialEq for Entry<'a> {
            fn eq(&self, other: &Self) -> bool {
                self.step() == other.step()
            }
        }
        impl<'a> Eq for Entry<'a> {}
        impl<'a> PartialOrd for Entry<'a> {
            fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
                other.step().partial_cmp(&self.step())
            }
        }
        impl<'a> Ord for Entry<'a> {
            fn cmp(&self, other: &Self) -> Ordering {
                other.step().cmp(&self.step())
            }
        }

        let queue = self
            .trials
            .iter()
            .filter(|t| !t.evaluations.is_empty())
            .map(|trial| Entry { trial, index: 0 })
            .collect::<BinaryHeap<_>>();

        struct Iter<'a> {
            queue: BinaryHeap<Entry<'a>>,
        }
        impl<'a> Iterator for Iter<'a> {
            type Item = &'a EvaluationRecord;

            fn next(&mut self) -> Option<Self::Item> {
                if let Some(mut entry) = self.queue.pop() {
                    let item = &entry.trial.evaluations[entry.index];
                    entry.index += 1;
                    if entry.index < entry.trial.evaluations.len() {
                        self.queue.push(entry);
                    }
                    Some(item)
                } else {
                    None
                }
            }
        }

        Iter { queue }
    }

    pub fn best_value(&self) -> Option<f64> {
        let problem_steps = self.problem.spec.steps.last();
        self.trials
            .iter()
            .filter_map(|t| t.value(problem_steps))
            .map(OrderedFloat)
            .min()
            .map(|x| x.0)
    }

    pub fn auc(&self, start_step: u64) -> Option<f64> {
        let vars = self.problem.spec.values_domain.variables();
        if vars.len() != 1 {
            return None;
        }

        let mut global_min = vars[0].range().low();
        if !global_min.is_finite() {
            global_min = 0.0;
        }

        let problem_steps = self.problem.spec.steps.last();
        let mut trials = self
            .trials
            .iter()
            .filter_map(|t| {
                if let (Some(step), Some(value)) = (t.end_step(), t.value(problem_steps)) {
                    Some((step, value))
                } else {
                    None
                }
            })
            .collect::<Vec<_>>();
        trials.sort_by_key(|t| t.0);

        let mut prev_step = 0;
        let mut current_min = std::f64::INFINITY;
        let mut auc = 0.0;
        for (mut step, value) in trials {
            if step <= start_step {
                step = start_step;
            } else {
                if prev_step == 0 {
                    return None;
                }
                auc += (current_min - global_min) * (step - prev_step) as f64;
            }

            if value < current_min {
                current_min = value;
            }
            prev_step = step;
        }

        let study_steps = self.budget * problem_steps;
        auc += (current_min - global_min) * (study_steps - prev_step) as f64;

        Some(auc / problem_steps as f64)
    }

    pub fn solver_elapsed(&self) -> Duration {
        self.trials.iter().map(|t| t.solver_elapsed()).sum()
    }

    pub fn first_complete_trial(&self) -> Option<&TrialRecord> {
        let problem_steps = self.problem.spec.steps.last();
        self.trials
            .iter()
            .filter(|t| t.steps() != problem_steps)
            .min_by_key(|t| t.start_step())
    }
}
