use crate::ast_node::{ASTContext, Node, TokenSlice};
use crate::operations::Operation;
use crate::parser::{comma_separate, find_matching_parenthesis, ParseNode, ParseNodeFunctional};
use crate::syntax_error::SyntaxError;
use crate::tokens::Token;
use crate::types::{FunctionSignature, IJType};
use anyhow::{Error, Result};
use itertools::Itertools;
use std::rc::Rc;

use super::next_node_functional;
use super::utils::gather_operands;

struct FunctionChain {
    functions: Vec<Rc<Node>>,
}

fn _get_output_type_from_function(function: Rc<Node>) -> Result<IJType, SyntaxError> {
    function
        .output_type
        .extract_signature()
        .map(|sig| *sig.output)
        .ok_or(SyntaxError::NotPureFunction(format!("{:?}", function)))
}

fn _get_input_types_from_function(function: Rc<Node>) -> Result<Vec<IJType>, SyntaxError> {
    match &function.output_type {
        IJType::Function(function_signature) => Ok(function_signature.input.clone()),
        _ => Err(SyntaxError::NotPureFunction(format!("{:?}", function))),
    }
}

fn _is_function(node: Rc<Node>) -> bool {
    matches!(node.output_type, IJType::Function(_))
}

impl FunctionChain {
    fn new(functions: Vec<Rc<Node>>) -> Result<Self, SyntaxError> {
        let chain = Self { functions };
        chain.verify_consistency()?;
        Ok(chain)
    }

    fn empty() -> Self {
        Self { functions: vec![] }
    }
    fn get_output_number_type(&self) -> Result<Option<String>> {
        if self.len() == 0 {
            return Ok(None);
        }
        let mut number_type =
            _get_output_type_from_function(self.functions.last().unwrap().clone())?
                .extract_number_type()
                .unwrap_or(None);

        for function in self.functions.iter().rev().skip(1) {
            let number_input_type = _get_input_types_from_function(function.clone())?[0]
                .maybe_apply_number_type(&number_type)
                .extract_number_type()
                .unwrap_or(None);
            number_type = _get_output_type_from_function(function.clone())?
                .maybe_apply_number_type(&number_input_type)
                .extract_number_type()
                .unwrap_or(None);
        }
        Ok(number_type)
    }
    fn get_input_number_type(&self) -> Result<Vec<Option<String>>> {
        if self.len() == 0 {
            return Ok(vec![]);
        }
        let mut number_types =
            _get_input_types_from_function(self.functions.first().unwrap().clone())?
                .iter()
                .map(|input_type| input_type.extract_number_type().unwrap_or(None))
                .collect::<Vec<Option<String>>>();

        for function in self.functions.iter().skip(1) {
            let number_output_type = _get_output_type_from_function(function.clone())?
                .maybe_apply_number_type(&number_types[0])
                .extract_number_type();
            number_types = match number_output_type {
                Some(None) => _get_input_types_from_function(function.clone())?
                    .iter()
                    .map(|input_type| {
                        input_type
                            .maybe_apply_number_type(&number_types[0])
                            .extract_number_type()
                            .unwrap_or(None)
                    })
                    .collect::<Vec<Option<String>>>(),
                _ => _get_input_types_from_function(function.clone())?
                    .iter()
                    .map(|input_type| {
                        input_type
                            .maybe_apply_number_type(&number_types[0])
                            .extract_number_type()
                            .unwrap_or(None)
                    })
                    .collect::<Vec<Option<String>>>(),
            }
        }
        Ok(number_types)
    }

    fn get_output_type(&self) -> Result<IJType> {
        let first_function = self.functions.first().ok_or_else(|| {
            SyntaxError::FunctionChainInconsistency("Function chain is empty".to_string())
        })?;
        let output_type = _get_output_type_from_function(first_function.clone())?;
        let number_type = self.get_output_number_type()?;
        Ok(output_type.maybe_apply_number_type(&number_type))
    }

    fn get_input_types(&self) -> Result<Vec<IJType>> {
        let last_function = self.functions.last().ok_or_else(|| {
            SyntaxError::FunctionChainInconsistency("Function chain is empty".to_string())
        })?;
        let input_types = _get_input_types_from_function(last_function.clone())?;
        let number_types = self.get_input_number_type()?;
        Ok(input_types
            .iter()
            .zip(number_types.iter())
            .map(|(input_type, number_type)| input_type.maybe_apply_number_type(number_type))
            .collect())
    }

    fn get_node_types(&self) -> Vec<IJType> {
        self.functions
            .iter()
            .map(|node| node.output_type.clone())
            .collect()
    }

    fn get_chain_type(&self) -> Result<IJType> {
        let input_types = self.get_input_types()?;
        let output_types = self.get_output_type()?;
        Ok(IJType::Function(FunctionSignature::new(
            input_types,
            output_types,
        )))
    }

    fn verify_consistency(&self) -> Result<(), SyntaxError> {
        for function in self.functions.iter() {
            if !_is_function(function.clone()) {
                return Err(SyntaxError::TypeError(
                    "Function".to_string(),
                    format!("{:?}", function.output_type),
                ));
            }
        }
        for (prev_function, next_function) in self.functions.iter().tuple_windows() {
            let prev_input_types = _get_input_types_from_function(prev_function.clone())?;
            if prev_input_types.len() != 1 {
                return Err(SyntaxError::FunctionChainInconsistency(format!(
                    "Function is not last in chain but does not have exactly one input type. Instead it has {:?}",
                    prev_input_types
                ))
                );
            }
            let prev_input_type = prev_input_types[0].clone();
            let next_output_type = _get_output_type_from_function(next_function.clone())?;
            if !prev_input_type.type_match(&next_output_type) {
                return Err(SyntaxError::FunctionChainInconsistency(format!(
                    "Function chain is not consistent: {:?} -> {:?}",
                    prev_input_type, next_output_type
                )));
            }
        }
        Ok(())
    }

    fn len(&self) -> usize {
        self.functions.len()
    }

    fn extend(&self, node: Rc<Node>) -> Result<Option<Self>> {
        if !_is_function(node.clone()) {
            return Err(SyntaxError::NotPureFunction(format!("{:?}", node)).into());
        }
        if self.len() == 0 {
            return Ok(Some(Self::new(vec![node])?));
        }

        let last_input_types = self.get_input_types()?;
        if last_input_types.len() != 1 {
            return Ok(None);
        }
        let last_input_type = last_input_types[0].clone();

        let node_output_type = _get_output_type_from_function(node.clone())?;
        if !last_input_type.type_match(&node_output_type) {
            return Ok(None);
        }

        let mut new_functions = self.functions.clone();
        new_functions.push(node);
        Ok(Some(Self::new(new_functions)?))
    }

    fn match_input_type(&self, input_types: Vec<IJType>) -> Result<bool> {
        let chain_input_types = self.get_input_types()?;
        Ok(chain_input_types.len() == input_types.len()
            && chain_input_types
                .iter()
                .zip(input_types.iter())
                .all(|(a, b)| a.type_match(b)))
    }
}

fn extend_chains(
    chains: Vec<FunctionChain>,
    nodes: Vec<Rc<Node>>,
) -> Result<Vec<FunctionChain>, SyntaxError> {
    let mut new_chains = vec![];
    new_chains.extend(chains.iter().flat_map(|chain| {
        nodes
            .iter()
            .filter_map(move |node| chain.extend(node.clone()).ok().flatten())
    }));
    Ok(new_chains)
}

pub struct FunctionComposition;
impl FunctionComposition {
    fn get_chains_from_slice(
        slice: TokenSlice,
        context: &mut ASTContext,
    ) -> Result<(Vec<FunctionChain>, TokenSlice)> {
        let next_token = context.get_token_at_index(slice.start)?;
        if next_token != &Token::LParen {
            return Err(SyntaxError::ExpectedToken(
                "(".to_string(),
                format!("found {:?}", next_token),
            )
            .into());
        }
        let slice = slice.move_start(1)?;
        let rparen_index =
            find_matching_parenthesis(context, slice, &Token::LParen, &Token::RParen)
                .map_err(|e| context.add_context_to_syntax_error(e, slice))?;
        let remainder = slice.move_start(rparen_index + 1)?;
        let slice = slice
            .move_end(rparen_index)
            .map_err(|e| context.add_context_to_syntax_error(e.into(), slice))?;
        if slice.is_empty() {
            return Err(context
                .add_context_to_syntax_error(SyntaxError::EmptyFunctionComposition.into(), slice));
        }

        let slices = comma_separate(slice, context)?;
        let function_operands = slices
            .into_iter()
            .map(|s| {
                let (nodes, rest) = next_node_functional(s, context, None)
                    .map_err(|e| context.add_context_to_syntax_error(e, s))?;
                if !rest.is_empty() {
                    return Err(context.add_context_to_syntax_error(
                        SyntaxError::UnhandledTokens(format!("found {:?}", rest)).into(),
                        rest,
                    ));
                }
                Ok(nodes)
            })
            .collect::<Result<Vec<Vec<Rc<Node>>>, Error>>()?;
        let chains = function_operands
            .into_iter()
            .try_fold(vec![FunctionChain::empty()], extend_chains)?;
        Ok((chains, remainder))
    }
}

impl ParseNode for FunctionComposition {
    fn next_node(
        _op: Token,
        slice: TokenSlice,
        context: &mut ASTContext,
    ) -> Result<(Rc<Node>, TokenSlice)> {
        let (chains, remainder) = Self::get_chains_from_slice(slice, context)?;
        let types = chains
            .iter()
            .map(|chain| chain.get_input_types())
            .collect::<Result<Vec<Vec<IJType>>>>()?;
        if types.is_empty() {
            return Err(context.add_context_to_syntax_error(
                SyntaxError::FunctionChainInconsistency(
                    "Types of function are not composible".to_string(),
                )
                .into(),
                slice,
            ));
        }
        let (value_operands, rest) = gather_operands(types, remainder, context)?;

        let operand_types = value_operands
            .iter()
            .map(|operand| operand.output_type.clone())
            .collect::<Vec<IJType>>();

        let mut matching_chain: Option<FunctionChain> = None;
        for chain in chains {
            if chain.match_input_type(operand_types.clone())? {
                if matching_chain.is_some() {
                    return Err(SyntaxError::FunctionChainAmbiguousType(format!(
                        "{:?}",
                        operand_types
                    ))
                    .into());
                }
                matching_chain = Some(chain);
            }
        }
        if matching_chain.is_none() {
            return Err(
                SyntaxError::FunctionChainAmbiguousType(format!("{:?}", operand_types)).into(),
            );
        }
        let matching_chain = matching_chain.unwrap();
        let mut full_input_type: Vec<IJType> = vec![];
        let mut all_nodes: Vec<Rc<Node>> = vec![];
        for node in matching_chain.functions.iter() {
            full_input_type.push(node.output_type.clone());
            all_nodes.push(node.clone());
        }
        full_input_type.extend(operand_types);
        all_nodes.extend(value_operands);

        let output_type = matching_chain.get_output_type()?;

        Ok((
            Rc::new(Node::new(
                Operation::FunctionComposition(matching_chain.len()),
                full_input_type,
                output_type,
                all_nodes,
                context,
            )?),
            rest,
        ))
    }
}

impl ParseNodeFunctional for FunctionComposition {
    fn next_node_functional_impl(
        _op: Token,
        slice: TokenSlice,
        context: &mut ASTContext,
        needed_outputs: Option<&[IJType]>,
    ) -> Result<(Vec<Rc<Node>>, TokenSlice)> {
        let slice = slice.move_start(1)?;
        let (chains, remainder) = Self::get_chains_from_slice(slice, context)?;

        let matching_chains = match needed_outputs {
            Some(needed_outputs) => chains
                .into_iter()
                .filter(|chain| {
                    needed_outputs
                        .iter()
                        .any(|output_type| match chain.get_output_type() {
                            Ok(chain_output_type) => chain_output_type == output_type.clone(),
                            Err(_) => false,
                        })
                })
                .collect(),
            None => chains,
        };

        let nodes: Vec<Rc<Node>> = matching_chains
            .into_iter()
            .map(|chain| {
                let chain_type = chain.get_chain_type()?;
                if let IJType::Function(_) = chain_type.clone() {
                    Ok(Rc::new(Node::new(
                        Operation::FunctionComposition(chain.len()),
                        chain.get_node_types(),
                        chain_type,
                        chain.functions.clone(),
                        context,
                    )?))
                } else {
                    Err(
                        SyntaxError::TypeError("Function".to_string(), format!("{:?}", chain_type))
                            .into(),
                    )
                }
            })
            .collect::<Result<Vec<Rc<Node>>>>()?;

        Ok((nodes, remainder))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ast_node::Node;
    use crate::function_enums::BinaryMathFunctionEnum;
    use crate::operations::Operation;
    use crate::parser::{parse_str, parse_str_no_context};
    use crate::types::FunctionSignature;

    fn _create_function_node(
        input_types: Vec<IJType>,
        output_type: IJType,
        context: &mut ASTContext,
    ) -> Result<Rc<Node>> {
        let function_signature = FunctionSignature::new(input_types, output_type);
        Ok(Rc::new(Node::new(
            Operation::Function("test_function".to_string()),
            vec![],
            IJType::Function(function_signature),
            vec![],
            context,
        )?))
    }

    fn _create_non_function_node(context: &mut ASTContext) -> Result<Rc<Node>> {
        Ok(Rc::new(Node::new(
            Operation::Symbol("test_symbol".to_string()),
            vec![],
            IJType::Tensor(None),
            vec![],
            context,
        )?))
    }

    #[test]
    fn test_function_chain_empty() {
        let chain = FunctionChain::empty();
        assert_eq!(chain.len(), 0);
    }

    #[test]
    fn test_function_chain_extend1() -> Result<()> {
        let mut context = ASTContext::new();
        let chain = FunctionChain::empty();
        let node = _create_function_node(
            vec![IJType::Tensor(None)],
            IJType::Tensor(None),
            &mut context,
        )?;
        let result = chain.extend(node);
        assert!(result.is_ok());
        let new_chain = result.unwrap().unwrap();
        assert_eq!(new_chain.len(), 1);

        Ok(())
    }

    #[test]
    fn test_function_chain_extend2() -> Result<()> {
        let mut context = ASTContext::new();
        let chain = FunctionChain::empty();
        let node1 = _create_function_node(
            vec![IJType::Tensor(None)],
            IJType::Tensor(None),
            &mut context,
        )?;
        let node2 = _create_function_node(
            vec![IJType::Tensor(None)],
            IJType::Tensor(None),
            &mut context,
        )?;
        let result = chain.extend(node1);
        assert!(result.is_ok());
        let new_chain = result.unwrap().unwrap();
        let result = new_chain.extend(node2);
        assert!(result.is_ok());
        let new_chain = result.unwrap().unwrap();
        assert_eq!(new_chain.len(), 2);

        Ok(())
    }

    #[test]
    fn test_function_chain_extend_incompatible() -> Result<()> {
        let mut context = ASTContext::new();
        let chain = FunctionChain::empty();
        let node1 = _create_function_node(
            vec![IJType::Tensor(None)],
            IJType::Tensor(None),
            &mut context,
        )?;
        let node2 = _create_function_node(
            vec![IJType::Tensor(None)],
            IJType::Number(None),
            &mut context,
        )?;
        let result = chain.extend(node1);
        assert!(result.is_ok());
        let new_chain = result.unwrap().unwrap();
        let result = new_chain.extend(node2);
        assert!(result.is_ok());
        let new_chain = result.unwrap();
        assert!(new_chain.is_none());

        Ok(())
    }

    #[test]
    fn test_function_chain_extend_not_function() -> Result<()> {
        let mut context = ASTContext::new();
        let chain = FunctionChain::empty();
        let node1 = _create_non_function_node(&mut context)?;
        let result = chain.extend(node1);
        assert!(result.is_err());

        Ok(())
    }

    #[test]
    fn test_function_composition_parser_minus_plus_scalar_scalar() {
        let result = parse_str_no_context("@(-,+) 1 2");
        assert!(result.is_ok());
        let (composition_node, _) = result.unwrap();
        let operands = composition_node.operands.clone();
        assert_eq!(operands.len(), 4);
        let node1 = operands[0].clone();
        assert_eq!(
            node1.output_type,
            IJType::Function(FunctionSignature::new(
                vec![IJType::Number(None)],
                IJType::Number(None),
            ))
        );
        let node2 = operands[1].clone();
        assert_eq!(node2.output_type, IJType::number_function(2));
        let node3 = operands[2].clone();
        assert_eq!(node3.output_type, IJType::Number(None));
        let node4 = operands[3].clone();
        assert_eq!(node4.output_type, IJType::Number(None));
    }

    #[test]
    fn test_function_composition_parser_minus_plus_scalar_tensor() {
        let result = parse_str_no_context("@(-,+) 1 [2]");
        assert!(result.is_ok());
        let (composition_node, _) = result.unwrap();
        let operands = composition_node.operands.clone();
        assert_eq!(operands.len(), 4);
        let node1 = operands[0].clone();
        assert_eq!(node1.output_type, IJType::tensor_function(1));
        let node2 = operands[1].clone();
        assert_eq!(
            node2.output_type,
            IJType::Function(FunctionSignature::new(
                vec![IJType::Number(None), IJType::Tensor(None)],
                IJType::Tensor(None),
            ))
        );
        let node3 = operands[2].clone();
        assert_eq!(node3.output_type, IJType::Number(None));
        let node4 = operands[3].clone();
        assert_eq!(node4.output_type, IJType::Tensor(None));
    }
    #[test]
    fn test_function_composition_parser_minus_plus_tensor_tensor() {
        let result = parse_str_no_context("@(-,+) [1] [2]");
        assert!(result.is_ok());
        let (composition_node, _) = result.unwrap();
        let operands = composition_node.operands.clone();
        assert_eq!(operands.len(), 4);
        let node1 = operands[0].clone();
        assert_eq!(node1.output_type, IJType::tensor_function(1));
        let node2 = operands[1].clone();
        assert_eq!(node2.output_type, IJType::tensor_function(2));
        let node3 = operands[2].clone();
        assert_eq!(node3.output_type, IJType::Tensor(None));
        let node4 = operands[3].clone();
        assert_eq!(node4.output_type, IJType::Tensor(None));
    }

    #[test]
    fn test_function_composition_parser_minus_tensor_scalar() {
        let result = parse_str_no_context("@(-) [1] 2");
        assert!(result.is_ok());
        let (composition_node, _) = result.unwrap();
        let operands = composition_node.operands.clone();
        assert_eq!(operands.len(), 3);
        let node1 = operands[0].clone();
        assert_eq!(
            node1.output_type,
            IJType::Function(FunctionSignature::new(
                vec![IJType::Tensor(None), IJType::Number(None)],
                IJType::Tensor(None),
            ))
        );
        let node2 = operands[1].clone();
        assert_eq!(node2.output_type, IJType::Tensor(None));
        let node3 = operands[2].clone();
        assert_eq!(node3.output_type, IJType::Number(None));
    }

    #[test]
    fn test_function_composition_parser_custom_function() {
        let mut context = ASTContext::new();
        parse_str("var f: Fn(N,N->N)", &mut context).unwrap();
        let result = parse_str("@(-, f) 1 2", &mut context);
        assert!(result.is_ok());
        let composition_node = result.unwrap();
        let operands = composition_node.operands.clone();
        assert_eq!(operands.len(), 4);
        let node1 = operands[0].clone();
        assert_eq!(node1.output_type, IJType::number_function(1));
        let node2 = operands[1].clone();
        assert_eq!(node2.output_type, IJType::number_function(2));
        let node3 = operands[2].clone();
        assert_eq!(node3.output_type, IJType::Number(None));
        let node4 = operands[3].clone();
        assert_eq!(node4.output_type, IJType::Number(None));
    }

    #[test]
    fn test_function_composition_empty() {
        assert!(parse_str_no_context("@()").is_err());
    }

    #[test]
    fn test_function_composition_functional() -> Result<()> {
        let (node, _) = parse_str_no_context("~@(+): Fn(N,N->N)")?;
        assert_eq!(node.op, Operation::FunctionComposition(1));

        let (node, _) = parse_str_no_context("~@(-,+): Fn(N,N->N)")?;
        assert_eq!(node.op, Operation::FunctionComposition(2));
        Ok(())
    }

    #[test]
    fn test_function_composition_nested() -> Result<()> {
        let (node, _) = parse_str_no_context("~@(@(-,-),+): Fn(N,N->N)")?;
        assert_eq!(node.op, Operation::FunctionComposition(2));
        let child1 = node.operands[0].clone();
        assert_eq!(child1.op, Operation::FunctionComposition(2));
        let child2 = node.operands[1].clone();
        assert_eq!(child2.op, Operation::BinaryFunction(BinaryMathFunctionEnum::Add));
        Ok(())
    }

    #[test]
    fn test_function_composition_number_type() -> Result<()> {
        let mut context = ASTContext::new();
        parse_str("var f: Fn(N<a>,N<b>->N<c>)", &mut context)?;
        parse_str("var g: Fn(N<c>->N<d>)", &mut context)?;
        let node = parse_str("@(g,f) 1 2", &mut context)?;
        assert_eq!(node.output_type, IJType::Number(Some("d".to_string())));

        parse_str("var f: Fn(N<a>,N<b>->N<c>)", &mut context)?;
        parse_str("var g: Fn(N->N)", &mut context)?;
        let node = parse_str("@(g,f) 1 2", &mut context)?;
        assert_eq!(node.output_type, IJType::Number(Some("c".to_string())));

        parse_str("var f: Fn(N<a>->N<a>)", &mut context)?;
        parse_str("var g: Fn(N->N<a>)", &mut context)?;
        let node = parse_str("@(g,f) 1", &mut context)?;
        assert_eq!(node.output_type, IJType::Number(Some("a".to_string())));
        let node = parse_str("@(f,g) 1", &mut context)?;
        assert_eq!(node.output_type, IJType::Number(Some("a".to_string())));

        parse_str("var f: Fn(N<a>->N<b>)", &mut context)?;
        parse_str("var g: Fn(N->N<a>)", &mut context)?;
        let node = parse_str("@(f,g) 1", &mut context)?;
        assert_eq!(node.output_type, IJType::Number(Some("b".to_string())));

        parse_str("var f: Fn(N<a>->N<b>)", &mut context)?;
        parse_str("var g: Fn(N->N<a>)", &mut context)?;
        let node = parse_str("@(f,g) 1", &mut context)?;
        assert_eq!(node.output_type, IJType::Number(Some("b".to_string())));
        let result = parse_str("@(f,g) 1<c>", &mut context);
        assert!(result.is_err());
        Ok(())
    }

    #[test]
    fn test_function_composition_function_type() -> Result<()> {
        let mut context = ASTContext::new();
        parse_str("var g: Fn(Fn(N<b>->N<c>)->N<d>)", &mut context)?;
        parse_str("var f: Fn(N<a> -> Fn(N<b>->N<c>))", &mut context)?;
        let node = parse_str("@(g,f) 1<a>", &mut context)?;
        assert_eq!(node.output_type, IJType::Number(Some("d".to_string())));

        let mut context = ASTContext::new();
        parse_str("var g: Fn(Fn(N->N)->N<d>)", &mut context)?;
        parse_str("var f: Fn(N<a> -> Fn(N<b>->N<c>))", &mut context)?;
        let result = parse_str("@(g,f) 1<a>", &mut context);
        assert!(result.is_err());

        Ok(())
    }

    #[test]
    fn test_function_composition_inconsistent() -> Result<()> {
        let mut context = ASTContext::new();
        parse_str("var g: Fn(N<a> -> N<b>)", &mut context)?;
        parse_str("var f: Fn(N<b> -> N<c>)", &mut context)?;
        let result = parse_str("@(g,f) 1<b>", &mut context);
        assert!(result.is_err());

        let mut context = ASTContext::new();
        parse_str("var g: Fn(N<b> -> N<a>)", &mut context)?;
        parse_str("var f: Fn(N -> N)", &mut context)?;
        let result = parse_str("@(g,f) 1<a>", &mut context);
        assert!(result.is_err());

        Ok(())
    }
}
