// Copyright 2018 Stefan Kroboth
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

#![recursion_limit = "512"]

extern crate proc_macro;
extern crate syn;
#[macro_use]
extern crate quote;

use proc_macro::TokenStream;
use syn::*;

#[proc_macro_derive(ArgminSolver)]
pub fn argminsolver(input: TokenStream) -> TokenStream {
    // parse the input tokens into a syntax tree
    let input: DeriveInput = syn::parse(input).unwrap();

    let name = &input.ident;
    let gen = &input.generics.clone();
    let whe = &input.generics.where_clause;
    let attrs = input.attrs.clone();

    let brackets: &[_] = &['(', ')'];
    let quotes: &[_] = &['"', '"'];
    let mut conditions: Vec<Expr> = vec![];
    let mut logs_str: Vec<String> = vec![];
    let mut logs_expr: Vec<Expr> = vec![];
    let mut solver_name = name.to_string();
    let mut tts;
    for attr in attrs.iter() {
        tts = &attr.tts;
        let path = &attr.path;
        let path = &quote!(#path).to_string();
        if path == "solver" {
            let tts2 = quote!(#tts).to_string();
            let attr = tts2.trim_matches(brackets).trim();
            solver_name = attr.to_string();
        }
        if path == "stop" {
            let tts2 = quote!(#tts).to_string();
            let attr = tts2.trim_matches(brackets).trim();
            let stuff: Vec<&str> = attr.split("=>").map(|x| x.trim()).collect();
            let stop_condition = stuff[0].trim_matches(quotes).to_string();
            let stop_reason = stuff[1].to_string();
            let bla = &format!(
                "
                if {condition} {{
                    self.base.set_termination_reason(TerminationReason::{reason});
                    return TerminationReason::{reason};
                }}
                ",
                condition = stop_condition,
                reason = stop_reason
            );
            conditions.push(syn::parse_str(bla).unwrap());
        }
        if path == "log" {
            let tts2 = quote!(#tts).to_string();
            let attr = tts2.trim_matches(brackets).trim();
            let stuff: Vec<&str> = attr.split("=>").map(|x| x.trim()).collect();
            let text = stuff[0].trim_matches(quotes).to_string();
            logs_str.push(text);
            let expr = stuff[1].trim_matches(quotes).to_string();
            logs_expr.push(syn::parse_str(&expr).unwrap());
        }
    }

    let expanded = quote! {
        impl #gen ArgminSolver for #name #gen #whe {
            /// Run the optimization algorithm
            fn run(&mut self) -> Result<ArgminResult<Self::Parameters>, Error> {
                let total_time = std::time::Instant::now();

                // do the inital logging
                let logs = make_kv!("max_iters" => self.max_iters();
                                    #(#logs_str => #logs_expr;)*);
                self.base.log_info(#solver_name, &logs)?;

                use std::sync::atomic::{AtomicBool, Ordering};
                use std::sync::Arc;

                let running = Arc::new(AtomicBool::new(true));

                #[cfg(feature = "ctrlc")]
                {
                    // Set up the Ctrl-C handler
                    use ctrlc;
                    let r = running.clone();
                    ctrlc::set_handler(move || {
                        r.store(false, Ordering::SeqCst);
                    })?;
                }

                self.init()?;

                while running.load(Ordering::SeqCst) {
                    // check first if it has already terminated
                    // This should probably be solved better.
                    // First, check if it isn't already terminated. If it isn't, evaluate the
                    // stopping criteria. If `self.terminate()` is called without the checking
                    // whether it has terminated already, then it may overwrite a termination set
                    // within `next_iter()`!
                    if !self.base.terminated() {
                        self.terminate();
                    }
                    // Now check once more if the algorithm has terminated. If yes, then break.
                    if self.base.terminated() {
                        break;
                    }

                    // Start time measurement
                    let start = std::time::Instant::now();

                    // execute iteration
                    let mut data = self.next_iter()?;

                    // End time measurement
                    let duration = start.elapsed();

                    // Set new current parameter
                    self.base.set_cur_param(data.param())
                             .set_cur_cost(data.cost());

                    // check if parameters are the best so far
                    if data.cost() <= self.base.best_cost() {
                        self.base.set_best_param(data.param())
                                 .set_best_cost(data.cost());
                    }

                    // logging
                    let mut log = self.base.kv_for_iter();
                    if let Some(ref mut iter_log) = data.get_kv() {
                        iter_log.push("time", duration.as_secs() as f64 +
                                              duration.subsec_nanos() as f64 * 1e-9);
                        log.merge(&mut iter_log.clone());

                    }
                    self.base.log_iter(&log)?;

                    // Write to file or something
                    self.base.write(&self.base.cur_param())?;

                    // increment iteration number
                    self.base.increment_iter();
                }

                // in case it stopped prematurely and `termination_reason` is still `NotTerminated`,
                // someone must have pulled the handbrake
                if self.base.cur_iter() < self.base.max_iters() && !self.base.terminated() {
                    self.base.set_termination_reason(TerminationReason::Aborted);
                }

                self.base.set_total_time(total_time.elapsed());

                let kv = make_kv!(
                    "termination_reason" => self.base.termination_reason();
                    "total_time" => self.base.total_time().as_secs() as f64 +
                                    self.base.total_time().subsec_nanos() as f64 * 1e-9;
                );

                self.base.log_info(
                    &format!("Terminated: {reason}", reason = self.base.termination_reason_text(),),
                    &kv,
                )?;

                Ok(self.base.result())
            }

            /// Run the essential parts of the optimization algorithm (no logging, no Ctrl-C
            /// handling)
            fn run_fast(&mut self) -> Result<ArgminResult<Self::Parameters>, Error> {
                self.init()?;

                loop {
                    // check first if it has already terminated
                    if !self.base.terminated() {
                        self.terminate();
                    }
                    // Now check once more if the algorithm has terminated. If yes, then break.
                    if self.base.terminated() {
                        break;
                    }

                    let mut data = self.next_iter()?;

                    // increment iteration number
                    self.base.increment_iter();

                    // Set new current parameter
                    self.base.set_cur_param(data.param())
                             .set_cur_cost(data.cost());

                    // check if parameters are the best so far
                    if data.cost() < self.base.best_cost() {
                        self.base.set_best_param(data.param())
                                 .set_best_cost(data.cost());
                    }
                }


                let mut kv = ArgminKV::new();
                Ok(self.base.result())
            }

            /// Applies the cost function or operator to a parameter vector `param`.
            /// Returns an `Err` if `apply` of `ArgminOperator` is not implemented.
            fn apply(&mut self, param: &Self::Parameters) -> Result<Self::OperatorOutput, Error> {
                self.base.apply(param)
            }

            /// Computes the gradient at parameter `param`.
            /// Returns an `Err` if `gradient` of `ArgminOperator` is not implemented.
            fn gradient(&mut self, param: &Self::Parameters) -> Result<Self::Parameters, Error> {
                self.base.gradient(param)
            }

            /// Computes the Hessian at parameter `param`.
            /// Returns an `Err` if `hessian` of `ArgminOperator` is not implemented.
            fn hessian(&mut self, param: &Self::Parameters) -> Result<Self::Hessian, Error> {
                self.base.hessian(param)
            }

            /// Returns the current parameter vector.
            fn cur_param(&self) -> Self::Parameters {
                self.base.cur_param()
            }

            /// Returns the most recently stored gradient.
            fn cur_grad(&self) -> Self::Parameters {
                self.base.cur_grad()
            }

            /// Returns the most recently stored Hessian.
            fn cur_hessian(&self) -> Self::Hessian {
                self.base.cur_hessian()
            }

            /// Sets the current parameter to `param`.
            fn set_cur_param(&mut self, param: Self::Parameters) {
                self.base.set_cur_param(param);
            }

            /// Sets the current gradient to `grad`.
            fn set_cur_grad(&mut self, grad: Self::Parameters) {
                self.base.set_cur_grad(grad);
            }

            /// Sets the current Hessian to `hessian`.
            fn set_cur_hessian(&mut self, hessian: Self::Hessian) {
                self.base.set_cur_hessian(hessian);
            }

            /// Sets the best parameter vector to `param`.
            fn set_best_param(&mut self, param: Self::Parameters) {
                self.base.set_best_param(param);
            }

            /// Modify the parameter vector by calling the `modify` method of the trait
            /// `ArgminOperator`. Will return an `Err` if `modify` is not implemented.
            fn modify(&self, param: &Self::Parameters, factor: f64) -> Result<Self::Parameters, Error> {
                self.base.modify(param, factor)
            }

            /// Returns the result of the optimization.
            fn result(&self) -> ArgminResult<Self::Parameters> {
                self.base.result()
            }

            /// Sets the maximum number of iterations to `iters`.
            fn set_max_iters(&mut self, iters: u64) {
                self.base.set_max_iters(iters);
            }

            /// Returns the maximum number of iterations.
            fn max_iters(&self) ->  u64 {
                self.base.max_iters()
            }

            /// Increments the iteration counter.
            fn increment_iter(&mut self) {
                self.base.increment_iter();
            }

            /// Returns the current number of iterations.
            fn cur_iter(&self) -> u64 {
                self.base.cur_iter()
            }

            /// Returns the most recently stored cost function value.
            fn cur_cost(&self) -> f64 {
                self.base.cur_cost()
            }

            /// Sets the current cost function value to `cost`
            fn set_cur_cost(&mut self, cost: f64) {
                self.base.set_cur_cost(cost);
            }

            /// Returns the best cost function value obtained so far.
            fn best_cost(&self) -> f64 {
                self.base.best_cost()
            }

            /// Sets the best cost function value.
            fn set_best_cost(&mut self, cost: f64) {
                self.base.set_best_cost(cost);
            }

            /// Sets the target cost function value to `cost`. The optimization algorithm will be
            /// terminated when this limit is reached.
            fn set_target_cost(&mut self, cost: f64) {
                self.base.set_target_cost(cost);
            }

            /// Increments the counter for the computations of the cost function by 1.
            fn increment_cost_func_count(&mut self) {
                self.base.increment_cost_func_count();
            }

            /// Increases the counter for the computations of the cost function by `count`.
            fn increase_cost_func_count(&mut self, count: u64) {
                self.base.increase_cost_func_count(count);
            }

            /// Returns the current value of the counter for the computations of the cost function.
            fn cost_func_count(&self) -> u64 {
                self.base.cost_func_count()
            }

            /// Increments the counter for the computations of the gradient by 1.
            fn increment_grad_func_count(&mut self) {
                self.base.increment_grad_func_count();
            }

            /// Increases the counter for the computations of the gradient by `count`.
            fn increase_grad_func_count(&mut self, count: u64) {
                self.base.increase_grad_func_count(count);
            }

            /// Returns the current value of the counter for the computations of the gradient.
            fn grad_func_count(&self) -> u64 {
                self.base.grad_func_count()
            }

            /// Increments the counter for the computations of the Hessian by 1.
            fn increment_hessian_func_count(&mut self) {
                self.base.increment_hessian_func_count();
            }

            /// Increases the counter for the computations of the Hessian by `count`.
            fn increase_hessian_func_count(&mut self, count: u64) {
                self.base.increase_hessian_func_count(count);
            }

            /// Returns the current value of the counter for the computations of the Hessian.
            fn hessian_func_count(&self) -> u64 {
                self.base.hessian_func_count()
            }

            /// Attaches a logger which implements `ArgminLog` to the solver.
            fn add_logger(&mut self, logger: Box<ArgminLog>) {
                self.base.add_logger(logger);
            }

            /// Attaches a writer which implements `ArgminWrite` to the solver.
            fn add_writer(&mut self, writer: Box<ArgminWrite<Param = Self::Parameters>>) {
                self.base.add_writer(writer);
            }

            /// Sets the `TerminationReason`
            fn set_termination_reason(&mut self, reason: TerminationReason) {
                self.base.set_termination_reason(reason);
            }

            /// Checks whether any of the conditions to terminate is true and terminates the
            /// algorithm.
            fn terminate(&mut self) -> TerminationReason {
                if self.base.cur_iter() >= self.base.max_iters() {
                    self.set_termination_reason(TerminationReason::MaxItersReached);
                    return TerminationReason::MaxItersReached;
                }

                if self.base.cur_cost() <= self.base.target_cost() {
                    self.set_termination_reason(TerminationReason::TargetCostReached);
                    return TerminationReason::TargetCostReached;
                }

                #(#conditions)*
                self.set_termination_reason(TerminationReason::NotTerminated);
                TerminationReason::NotTerminated
            }

            /// Resets the `base` field to it's initial conditions. This is helpful for
            /// implementing a solver which is initialized once, but called several times. It is
            /// recommended to only call this method inside the `init` function of
            /// `ArgminNextIter`.
            fn base_reset(&mut self) {
                self.base.reset();
            }
        }
    };
    expanded.into()
}

#[proc_macro_derive(ArgminOperator)]
pub fn argminoperator(input: TokenStream) -> TokenStream {
    // parse the input tokens into a syntax tree
    let input: DeriveInput = syn::parse(input).unwrap();

    let name = &input.ident;
    let gen = &input.generics.clone();
    let whe = &input.generics.where_clause;
    let attrs = input.attrs.clone();

    let brackets: &[_] = &['(', ')'];
    // let quotes: &[_] = &['"', '"'];
    let mut parameters: Option<Type> = None;
    let mut output: Option<Type> = None;
    let mut hessian_type: Option<Type> = None;
    let mut cost_function: Option<Ident> = None;
    let mut gradient: Option<Item> = None;
    let mut hessian: Option<Item> = None;
    let mut modify: Option<Item> = None;
    let mut tts;
    for attr in attrs.iter() {
        tts = &attr.tts;
        let path = &attr.path;
        let path = &quote!(#path).to_string();
        if path == "parameters_type" {
            let tts2 = quote!(#tts).to_string();
            let attr = tts2.trim_matches(brackets).trim();
            parameters = Some(syn::parse_str(attr).unwrap());
        }
        if path == "output_type" {
            let tts2 = quote!(#tts).to_string();
            let attr = tts2.trim_matches(brackets).trim();
            output = Some(syn::parse_str(attr).unwrap());
        }
        if path == "hessian_type" {
            let tts2 = quote!(#tts).to_string();
            let attr = tts2.trim_matches(brackets).trim();
            hessian_type = Some(syn::parse_str(attr).unwrap());
        }
        if path == "cost_function" {
            let tts2 = quote!(#tts).to_string();
            let attr = tts2.trim_matches(brackets).trim();
            cost_function = Some(syn::parse_str(attr).unwrap());
        }
        if path == "gradient" {
            let tts2 = quote!(#tts).to_string();
            let attr = tts2.trim_matches(brackets).trim();
            let tmp = format!(
                "fn gradient(&self, x: &Self::Parameters) -> Result<Self::Parameters, Error> {{ 
                    Ok({}(x))
                 }}",
                attr
            );
            gradient = Some(syn::parse_str(&tmp).unwrap());
        }
        if path == "hessian" {
            let tts2 = quote!(#tts).to_string();
            let attr = tts2.trim_matches(brackets).trim();
            let tmp = format!(
                "fn hessian(&self, x: &Self::Parameters) -> Result<Self::Hessian, Error> {{ 
                    Ok({}(x))
                 }}",
                attr
            );
            hessian = Some(syn::parse_str(&tmp).unwrap());
        }
        if path == "modify" {
            let tts2 = quote!(#tts).to_string();
            let attr = tts2.trim_matches(brackets).trim();
            let tmp = format!(
                "fn modify(&self, x: &Self::Parameters, scale: f64) -> Result<Self::Parameters, Error> {{ 
                    Ok({}(x, scale))
                 }}",
                attr
            );
            modify = Some(syn::parse_str(&tmp).unwrap());
        }
    }

    let expanded = quote! {
        impl #gen ArgminOperator for #name #gen #whe {
            type Parameters = #parameters;
            type OperatorOutput = #output;
            type Hessian = #hessian_type;

            fn apply(&self, param: &Self::Parameters) ->  Result<Self::OperatorOutput, Error> {
                Ok(#cost_function(param))
            }

            #gradient

            #hessian

            #modify
        }
    };
    expanded.into()
}