use canvas::Draw2D;
use colored::*;
use gnuplot as gp;
use itertools_num::linspace;
use minuit::FitOutput;
use prelude::ParamInfo;
use std::{iter, marker::PhantomData, mem, ops::Fn, rc::Rc};

pub struct ParamIterator<'a> {
    p: &'a [f64],
    pinfo: &'a [ParamInfo],
    count: usize,
}

impl<'a> ParamIterator<'a> {
    pub fn new(p: &'a [f64], pinfo: &'a [ParamInfo]) -> Self {
        ParamIterator { p, pinfo, count: 0 }
    }
}

impl<'a> Iterator for ParamIterator<'a> {
    type Item = (&'a f64, &'a ParamInfo);
    fn next(&mut self) -> Option<Self::Item> {
        if self.count < self.p.len() {
            self.count += 1;
            Some((&self.p[self.count - 1], &self.pinfo[self.count - 1]))
        } else {
            None
        }
    }
}

pub struct ParamIteratorMut<'a> {
    p: &'a mut [f64],
    pinfo: &'a mut [ParamInfo],
    count: usize,
}

impl<'a> ParamIteratorMut<'a> {
    pub fn new(p: &'a mut [f64], pinfo: &'a mut [ParamInfo]) -> Self {
        ParamIteratorMut { p, pinfo, count: 0 }
    }
}

impl<'a> Iterator for ParamIteratorMut<'a> {
    type Item = (&'a mut f64, &'a mut ParamInfo);
    fn next(&mut self) -> Option<Self::Item> {
        if self.count < self.p.len() {
            self.count += 1;
            unsafe {
                mem::transmute(Some((
                    &mut self.p[self.count - 1],
                    &mut self.pinfo[self.count - 1],
                )))
            }
        } else {
            None
        }
    }
}

#[derive(Clone)]
pub struct Function<'a, 'b, F: 'b + Fn(f64, &[f64]) -> f64> {
    name: Option<String>,
    func: Rc<F>,
    p: Vec<f64>,
    pinfo: Vec<ParamInfo>,
    fmt: Vec<gp::PlotOption<&'a str>>,
    xmin: f64,
    xmax: f64,
    divd: usize,
    _phantom: PhantomData<&'b u8>,
}

impl<'a, 'c, F> Function<'a, 'c, F>
where
    F: 'c + Fn(f64, &[f64]) -> f64,
{
    pub fn new(npar: usize, _f: F) -> Self {
        Function {
            name: None,
            func: Rc::new(_f),
            p: iter::repeat(0f64).take(npar).collect(),
            pinfo: iter::repeat(ParamInfo::new(1e-1f64, None))
                .take(npar)
                .collect(),
            fmt: vec![],
            xmin: 1f64,
            xmax: 0f64,
            divd: 1000,
            _phantom: PhantomData,
        }
    }

    pub fn title<'b>(&'b self) -> Option<&'b str> { self.name.as_ref().map(|t| t.as_str()) }

    pub fn with_title<S: AsRef<str>>(mut self, title: S) -> Self {
        self.set_title(title);
        self
    }

    pub fn set_title<S: AsRef<str>>(&mut self, title: S) {
        self.name = Some(String::from(title.as_ref()));
    }

    pub fn fmt<'b>(mut self, opts: &'b [gp::PlotOption<&'a str>]) -> Self {
        self.set_fmt(opts);
        self
    }

    pub fn set_fmt<'b>(&mut self, opts: &'b [gp::PlotOption<&'a str>]) {
        self.fmt = opts.iter().cloned().collect();
    }

    pub fn par(mut self, idx: usize, val: f64, err: f64, rng: Option<(f64, f64)>) -> Self {
        self.set_parameter(idx, val, err, rng);
        self
    }

    pub fn set_parameter<'b>(
        &'b mut self,
        idx: usize,
        val: f64,
        err: f64,
        rng: Option<(f64, f64)>,
    ) {
        self.p[idx] = val;
        self.pinfo[idx] = ParamInfo::new(err, rng);
    }

    pub fn pars<'b>(&'b self) -> ParamIterator<'b> { ParamIterator::new(&self.p, &self.pinfo) }

    pub fn pars_mut<'b>(&'b mut self) -> ParamIteratorMut<'b> {
        ParamIteratorMut::new(&mut self.p, &mut self.pinfo)
    }

    pub fn xrange(mut self, left: Option<f64>, right: Option<f64>) -> Self {
        self.set_xrange(left, right);
        self
    }

    pub fn set_xrange(&mut self, left: Option<f64>, right: Option<f64>) {
        if let Some(x) = left {
            self.xmin = x;
        }
        if let Some(x) = right {
            self.xmax = x;
        }
    }

    pub fn print_pars(&self) -> String {
        let mut resstr = String::new();
        for (i, (val, meta)) in self.p.iter().zip(self.pinfo.iter()).enumerate() {
            resstr += format!("p{}: {:<9.5E} +/- {:<9.5E}\n", i, val, meta.err).as_str();
        }
        resstr
    }

    pub fn func(&self) -> &Rc<F> { &self.func }
}

impl<'a, 'f, 'c, F> From<(&'c Function<'a, 'f, F>, FitOutput)> for Function<'a, 'f, F>
where
    F: 'f + Fn(f64, &[f64]) -> f64,
{
    fn from(fo: (&Function<'a, 'f, F>, FitOutput)) -> Self {
        Function {
            p: (fo.1).0,
            pinfo: (fo.1).1,
            name: None,
            func: fo.0.func().clone(),
            fmt: vec![],
            ..*fo.0
        }
    }
}

impl<'a, 'b, F> FnOnce<(f64,)> for Function<'a, 'b, F>
where
    F: 'b + Fn(f64, &[f64]) -> f64,
{
    type Output = f64;
    extern "rust-call" fn call_once(self, (x): (f64,)) -> Self::Output {
        (self.func)(x.0, self.p.as_slice())
    }
}

impl<'a, 'b, F> FnMut<(f64,)> for Function<'a, 'b, F>
where
    F: 'b + Fn(f64, &[f64]) -> f64,
{
    extern "rust-call" fn call_mut(&mut self, (x): (f64,)) -> f64 {
        (self.func)(x.0, self.p.as_slice())
    }
}

impl<'a, 'b, F> Fn<(f64,)> for Function<'a, 'b, F>
where
    F: 'b + Fn(f64, &[f64]) -> f64,
{
    extern "rust-call" fn call(&self, (x): (f64,)) -> f64 { (self.func)(x.0, self.p.as_slice()) }
}

impl<'a, 'b, F> Draw2D for Function<'a, 'b, F>
where
    F: 'b + Fn(f64, &[f64]) -> f64,
{
    fn draw2d(&self, _axes: &mut gp::Axes2D) {
        if self.xmin < self.xmax {
            let x = linspace::<f64>(self.xmin, self.xmax, self.divd);
            let y = x.clone().map(|x| (self.func)(x, self.p.as_slice()));
            _axes.lines(
                x,
                y,
                self.fmt
                    .iter()
                    .cloned()
                    .chain(self.name.as_ref().map(|t| gp::Caption(t.as_str())))
                    .collect::<Vec<_>>()
                    .as_slice(),
            );
        } else {
            println!(
                "{} > {}: uncpecified limits.",
                self.name.as_ref().map_or("fit", |t| t.as_str()).magenta(),
                "Warning".yellow()
            );
        }
    }
}