#include "externals.hh"

#include "c-opopt.hh"
#include "c-ops.hh"
#include "c-sync.hh"

using namespace ebcl;
using namespace opast;
using namespace opopt;


/*= T_OptData ================================================================*/

void T_OptData::findInputDecls(
		T_OpsParserOutput& program ) noexcept
{
	if ( inputDecls ) {
		return;
	}

	inputDecls = T_KeyValueTable< T_String , T_Array< T_InputDecl > >{ };
	visitor.visit( program.root , [this]( A_Node& node , const bool exit ) {
		if ( exit && node.type( ) == A_Node::OP_INPUT ) {
			auto& input{ (T_InputInstrNode&) node };
			auto* da{ inputDecls->get( input.id( ) ) };
			if ( !da ) {
				inputDecls->add( input.id( ) , T_Array< T_InputDecl >{ } );
				da = inputDecls->get( input.id( ) );
			}
			da->add( T_InputDecl{ input.location( ) , input.defValue( ) } );
		}
		return true;
	} );
}


/*= CONSTANT FOLDING =========================================================*/

namespace {

struct T_ConstantFolder_
{
	T_ConstantFolder_( T_OptData& data ) noexcept
		: oData{ data }
	{}

	// Result
	bool didFold{ false };

	bool operator()( A_Node& node , bool exit ) noexcept;

    private:
	T_OptData& oData;

	template<
		typename T
	> void handleParentNode(
			A_Node& node ,
			std::function< A_ExpressionNode&( T& ) > get ,
			std::function< void( T& , P_ExpressionNode ) > set ) noexcept;

	P_ExpressionNode checkExpression(
			A_ExpressionNode& node ) noexcept;

	// Handle identifiers. If the size is fixed and the identifier is
	// either width or height, replace it with the appropriate value.
	P_ExpressionNode doIdExpr(
			T_IdentifierExprNode& node ) noexcept;

	// Handle reads from inputs. If there's a curve and it is a constant,
	// or if there's no curve and only one default value, then the
	// expression is constant.
	P_ExpressionNode doInputExpr(
			T_InputExprNode& node ) noexcept;

	// Transform an unary operator applied to a constant into a constant.
	P_ExpressionNode doUnaryOp(
			T_UnaryOperatorNode& node ,
			double value ) const noexcept;

	// Transform a binary operator applied to a constant into a constant.
	P_ExpressionNode doBinaryOp(
			T_BinaryOperatorNode& node ,
			double left ,
			double right ) const noexcept;
};

/*----------------------------------------------------------------------------*/

bool T_ConstantFolder_::operator()(
		A_Node& node ,
		const bool exit ) noexcept
{
	if ( exit ) {
		return true;
	}

	switch ( node.type( ) ) {

	    case A_Node::TN_ARG:
		handleParentNode< T_ArgumentNode >(
			node ,
			[]( auto& n ) -> A_ExpressionNode& { return n.expression( ); } ,
			[]( auto& n , P_ExpressionNode e ) { n.expression( std::move( e ) ); }
		);
		return false;

	    case A_Node::TN_CONDITION:
		handleParentNode< T_CondInstrNode::T_Expression >( node ,
			[]( auto& n ) -> A_ExpressionNode& { return n.expression( ); } ,
			[]( auto& n , P_ExpressionNode e ) { n.expression( std::move( e ) ); }
		);
		return false;

	    case A_Node::OP_SET:
		handleParentNode< T_SetInstrNode >( node ,
			[]( auto& n ) -> A_ExpressionNode& { return n.expression( ); } ,
			[]( auto& n , P_ExpressionNode e ) { n.setExpression( std::move( e ) ); } );
		return false;

	    default:
		return true;
	}
}

/*----------------------------------------------------------------------------*/

template<
	typename T
> void T_ConstantFolder_::handleParentNode(
		A_Node& n ,
		std::function< A_ExpressionNode&( T& ) > get ,
		std::function< void( T& , P_ExpressionNode ) > set ) noexcept
{
	auto& node{ (T&) n };
	auto& child{ get( node ) };
	auto r{ checkExpression( child ) };
	if ( r ) {
		oData.logger( [&]() {
			T_StringBuilder sb;
			sb << "substituting node at " << child.location( );
			return sb;
		} , 2 );
		r->location( ) = node.location( );
		set( node , std::move( r ) );
		didFold = true;
	}
}

P_ExpressionNode T_ConstantFolder_::checkExpression(
		A_ExpressionNode& node ) noexcept
{
	// Already a constant
	if ( node.type( ) == A_Node::EXPR_CONST ) {
		return {};
	}

	// Replace $width/$height with value if fixedSize
	if ( node.type( ) == A_Node::EXPR_ID ) {
		return doIdExpr( (T_IdentifierExprNode&) node );
	}

	// Replace inputs with value if no curve/constant curve
	if ( node.type( ) == A_Node::EXPR_INPUT ) {
		return doInputExpr( (T_InputExprNode&) node );
	}

	// Replace UnOp( Cnst ) with result
	auto* const asUnary{ dynamic_cast< T_UnaryOperatorNode* >( &node ) };
	if ( asUnary ) {
		handleParentNode< T_UnaryOperatorNode >( *asUnary ,
			[]( auto& n ) -> A_ExpressionNode& { return n.argument( ); } ,
			[]( auto& n , P_ExpressionNode e ) { n.setArgument( std::move( e ) ); } );
		if ( asUnary->argument( ).type( ) == A_Node::EXPR_CONST ) {
			auto const& cn{ (T_ConstantExprNode const&) asUnary->argument( ) };
			return doUnaryOp( *asUnary , cn.floatValue( ) );
		}
		return {};
	}

	// Replace BinOp( Cnst , Cnst ) with result
	auto* const asBinary{ dynamic_cast< T_BinaryOperatorNode* >( &node ) };
	assert( asBinary && "Missing support for some expr subtype" );
	handleParentNode< T_BinaryOperatorNode >( *asBinary ,
		[]( auto& n ) -> A_ExpressionNode& { return n.left( ); } ,
		[]( auto& n , P_ExpressionNode e ) { n.setLeft( std::move( e ) ); } );
	handleParentNode< T_BinaryOperatorNode >( *asBinary ,
		[]( auto& n ) -> A_ExpressionNode& { return n.right( ); } ,
		[]( auto& n , P_ExpressionNode e ) { n.setRight( std::move( e ) ); } );

	if ( asBinary->left( ).type( ) == A_Node::EXPR_CONST
			&& asBinary->right( ).type( ) == A_Node::EXPR_CONST ) {
		auto const& l{ (T_ConstantExprNode const&) asBinary->left( ) };
		auto const& r{ (T_ConstantExprNode const&) asBinary->right( ) };
		return doBinaryOp( *asBinary , l.floatValue( ) , r.floatValue( ) );
	}
	return {};
}

/*----------------------------------------------------------------------------*/

P_ExpressionNode T_ConstantFolder_::doInputExpr(
		T_InputExprNode& node ) noexcept
{
	if ( !oData.curves ) {
		return {};
	}

	auto const* const curve{ oData.curves->curves.get( node.id( ) ) };
	if ( curve ) {
		// Curve present, check if it's constant
		const auto cval{ curve->isConstant( ) };
		if ( !cval ) {
			return {};
		}
		return NewOwned< T_ConstantExprNode >( node.parent( ) , *cval );
	}

	assert( oData.inputDecls );
	auto const* const dva{ oData.inputDecls->get( node.id( ) ) };
	assert( dva );
	if ( dva->size( ) == 1 ) {
		// If there's only one default value, that's a constant.
		return NewOwned< T_ConstantExprNode >( node.parent( ) ,
				(*dva)[ 0 ].value );
	}
	return {};
}

P_ExpressionNode T_ConstantFolder_::doIdExpr(
		T_IdentifierExprNode& node ) noexcept
{
	if ( !oData.fixedSize ) {
		return {};
	}

	if ( node.id( ) == "width" ) {
		oData.logger( []{
			return T_StringBuilder{ "replacing $width with fixed width" };
		} , 2 );
		return NewOwned< T_ConstantExprNode >( node.parent( ) ,
				double( oData.fixedSize->first ) );
	}

	if ( node.id( ) == "height" ) {
		oData.logger( []{
			return T_StringBuilder{ "replacing $height with fixed height" };
		} , 2 );
		return NewOwned< T_ConstantExprNode >( node.parent( ) ,
				float( oData.fixedSize->second ) );
	}

	return {};
}

P_ExpressionNode T_ConstantFolder_::doUnaryOp(
		T_UnaryOperatorNode& node ,
		const double value ) const noexcept
{
	const double rVal{ [this]( auto& node , const auto value ) {
		switch ( node.op( ) ) {

		    case T_UnaryOperatorNode::NEG:
			return -value;

		    case T_UnaryOperatorNode::NOT:
			return value ? 0. : 1.;

		    case T_UnaryOperatorNode::INV:
			if ( value == 0 ) {
				oData.errors.addNew( "math - 1/x, x=0" , node.location( ) );
				return 0.;
			}
			return 1. / value;

		    case T_UnaryOperatorNode::COS:
			return cos( value );

		    case T_UnaryOperatorNode::SIN:
			return sin( value );

		    case T_UnaryOperatorNode::TAN:
			if ( fabs( value - M_PI / 2 ) <= 1e-6 ) {
				oData.errors.addNew( "math - tan(x), x=~PI/2" ,
						node.location( ) , E_SRDErrorType::WARNING );
			}
			return tan( value );

		    case T_UnaryOperatorNode::SQRT:
			if ( value < 0 ) {
				oData.errors.addNew( "math - sqrt(x), x<0" , node.location( ) );
				return 0.;
			}
			return sqrt( value );

		    case T_UnaryOperatorNode::LN:
			if ( value <= 0 ) {
				oData.errors.addNew( "math - ln(x), x<=0" , node.location( ) );
				return 0.;
			}
			return log( value );

		    case T_UnaryOperatorNode::EXP:
			return exp( value );
		}

		fprintf( stderr , "invalid operator %d\n" , int( node.op( ) ) );
		std::abort( );
	}( node , value ) };

	return NewOwned< T_ConstantExprNode >( node.parent( ) , rVal );
}

P_ExpressionNode T_ConstantFolder_::doBinaryOp(
		T_BinaryOperatorNode& node ,
		const double left ,
		const double right ) const noexcept
{
	const double rVal{ [this]( auto& node , const auto l , const auto r ) {
		switch ( node.op( ) ) {

		    case T_BinaryOperatorNode::ADD:
			return l + r;

		    case T_BinaryOperatorNode::SUB:
			return l - r;

		    case T_BinaryOperatorNode::MUL:
			return l * r;

		    case T_BinaryOperatorNode::DIV:
			if ( r == 0 ) {
				oData.errors.addNew( "math - l/r, r=0" , node.location( ) );
				return 0.;
			}
			return l / r;

		    case T_BinaryOperatorNode::POW:
			if ( l == 0 && r == 0 ) {
				oData.errors.addNew( "math - l^r, l=r=0" , node.location( ) );
				return 0.;
			}
			if ( l == 0 && r < 0 ) {
				oData.errors.addNew( "math - l^r, l=0, r<0" , node.location( ) );
				return 0.;
			}
			if ( l < 0 && fmod( r , 1. ) != 0. ) {
				oData.errors.addNew( "math - l^r, l<0, r not integer" , node.location( ) );
				return 0.;
			}
			return pow( l , r );

		    case T_BinaryOperatorNode::CMP_EQ: return ( l == r ) ? 1. : 0.;
		    case T_BinaryOperatorNode::CMP_NE: return ( l != r ) ? 1. : 0.;
		    case T_BinaryOperatorNode::CMP_GT: return ( l >  r ) ? 1. : 0.;
		    case T_BinaryOperatorNode::CMP_GE: return ( l >= r ) ? 1. : 0.;
		    case T_BinaryOperatorNode::CMP_LT: return ( l <  r ) ? 1. : 0.;
		    case T_BinaryOperatorNode::CMP_LE: return ( l <= r ) ? 1. : 0.;
		}
		fprintf( stderr , "invalid operator %d\n" , int( node.op( ) ) );
		std::abort( );
	}( node , left , right ) };

	return NewOwned< T_ConstantExprNode >( node.parent( ) , rVal );
}

} // namespace <anon>

/*----------------------------------------------------------------------------*/


bool opopt::FoldConstants(
		T_OpsParserOutput& program ,
		T_OptData& oData ) noexcept
{
	T_ConstantFolder_ folder{ oData };
	oData.logger( []() {
		return T_StringBuilder{ "constant folding pass" };
	} , 1 );
	if ( oData.curves ) {
		oData.findInputDecls( program );
	}
	oData.visitor.visit( program.root , [&]( auto& n , auto x ) {
		return folder( n , x );
	} );
	oData.logger( [&]() {
		return T_StringBuilder{
			folder.didFold ? "some constants were folded"
				: "no constants were folded" };
	} , 1 );
	return folder.didFold;
}


/*= DEAD CODE REMOVAL ========================================================*/

bool opopt::RemoveDeadCode(
		T_OpsParserOutput& program ,
		T_OptData& oData ) noexcept
{
	return false;
}