#include <ginac/ginac.h>

namespace GiNaC {
   class vector;
   class Inner_Prod;
   class Cross_Prod;
   ex inner_prod(ex const& l, ex const& r);
   ex cross_prod(ex const& l, ex const& r);

   inline bool is_a_vector(ex const& e)
   {
      return is_a<lst>(e) || is_a<vector>(e) || is_a<Cross_Prod>(e);
   }

   std::pair<ex,ex> get_scalar_and_vector(ex const& e)
   {
      if(!is_a<mul>(e))
      {
         if(is_a_vector(e))
            return std::pair<ex, ex>(1, e);
         else
            return std::pair<ex, ex>(e, 1);
      }

      ex scalar = 1;
      ex vector;
      for(const_iterator itr = e.begin(); itr != e.end(); ++itr)
      {
         ex term = *itr;
         if(is_a_vector(term))
         {
            assert(vector.nops() == 0);
            vector = term;
         }
         else
            scalar *= term;
      }
      return std::pair<ex, ex>(scalar, vector);
   }

   template<class Product>
   ex cannonicalize_vector_product(ex const& lhs, ex const& rhs, Product prod)
   {
      std::pair<ex, ex> lhs_sv = get_scalar_and_vector(lhs);
      std::pair<ex, ex> rhs_sv = get_scalar_and_vector(rhs);

      ex scalar = lhs_sv.first * rhs_sv.first;
      if(lhs_sv.second == 1 || rhs_sv.second == 1)
         return scalar * lhs_sv.second * rhs_sv.second;

      return scalar * prod(lhs_sv.second, rhs_sv.second);
   }

   template<class Product>
   ex expand_vector_product(
      ex const& lhs, ex const& rhs, Product prod, unsigned options = 0)
   {
      ex l = lhs.expand(options);
      ex r = rhs.expand(options);
      ex res = 0;

      // tricky stuff here
      //    first we want to use iterator to traverse op sequence to get
      //    O(n) irrespective of how the tree is represented
      //    second traverse the iterator only if ex is an add sequence
      //    remember that nops() can be larger than one in a number of ways
      //    also that nops() is zero for an atom
      bool lhs_is_singleton = (!is_a<add>(l) || l.nops() == 0);
      int n_lhs_terms = (lhs_is_singleton ? 1 : l.nops());
      const_iterator litr;
      for(int li = 0; li < n_lhs_terms; ++li)
      {
         if(li == 0 && !lhs_is_singleton)
            litr = l.begin();
         ex le = (li == 0 && lhs_is_singleton ? l : *litr);
         bool rhs_is_singleton = (!is_a<add>(r) || r.nops() == 0);
         int n_rhs_terms = (rhs_is_singleton ? 1 : r.nops());
         const_iterator ritr;
         for(int ri = 0; ri < n_rhs_terms; ++ri)
         {
            if(ri == 0 && !rhs_is_singleton)
               ritr = r.begin();
            ex re = (ri == 0 && rhs_is_singleton ? r : *ritr);
            res += cannonicalize_vector_product(le, re, prod);
            if(!rhs_is_singleton)
               ++ritr;
         }
         if(!lhs_is_singleton)
            ++litr;
      }

      return res;
   }



   class component : public basic
   {
      GINAC_DECLARE_REGISTERED_CLASS(component, basic);

   public:
      ex  base;
      int index;

      component(ex e, int i)
       : inherited(&component::tinfo_static)
       , base(e)
       , index(i)
      {
      }   

      ex subs(exmap const&m, unsigned options = 0) const
      {
         ex subex = inherited::subs(m, options);
         if(!is_a<component>(subex))
            return subex;

         component sub_self = ex_to<component>(subex);
         if(!is_a<lst>(sub_self.base))
            return subex;

         if(sub_self.index >= sub_self.base.nops())
             throw std::range_error("component::op(): no such operand");

         return sub_self.base.op(sub_self.index);
      }

      size_t nops() const
      {
         return 1;
      }

      ex op(size_t i) const
      {
         switch (i) {
         case 0:
             return this->base;
         default:
             throw std::range_error("Inner_Prod::op(): no such operand");
         }
      }

      ex& let_op(size_t i)
      {
         ensure_if_modifiable();

         switch (i) {
         case 0:
             return this->base;
         default:
             throw std::range_error("Inner_Prod::op(): no such operand");
         }
      }

      void do_print(print_context const& c, unsigned level = 0) const;
   };

   GINAC_IMPLEMENT_REGISTERED_CLASS_OPT(
      component, basic, print_func<print_context>(&component::do_print));

   component::component()
    : inherited(&component::tinfo_static)
   {
   }

   component::component(archive_node const& n, lst& sym_lst)
    : inherited(n, sym_lst)
   {
   }

   void component::archive(archive_node& n) const
   {
      inherited::archive(n);
   }

   ex component::unarchive(archive_node const& n, lst& sym_lst)
   {
      return (new component(n, sym_lst))->setflag(status_flags::dynallocated);
   }

   int component::compare_same_type(basic const& other) const
   {
      if(!is_a<component>(other))
         return 1;

      component const& o = static_cast<component const&>(other);
      return (this->base == o.base && this->index == o.index ? 0 : 1);
   }

   void component::do_print(print_context const& c, unsigned level) const
   {
      // print_context::s is a reference to an ostream
      c.s << this->base << '[' << this->index << ']';
   }







   class vector : public symbol
   {
      GINAC_DECLARE_REGISTERED_CLASS(vector, symbol);

   public:
      vector(std::string const& s)
       : inherited(s)
      {
      }

      ex operator[](int i) const
      {
         return component(*this, i);
      }
   };

   GINAC_IMPLEMENT_REGISTERED_CLASS(vector, symbol);

   vector::vector()
   {
      tinfo_key = &vector::tinfo_static; 
   }

   vector::vector(archive_node const& n, lst& sym_lst)
    : inherited(n, sym_lst)
   {
   }

   void vector::archive(archive_node& n) const
   {
      inherited::archive(n);
   }

   ex vector::unarchive(archive_node const& n, lst& sym_lst)
   {
      return (new vector(n, sym_lst))->setflag(status_flags::dynallocated);
   }

   int vector::compare_same_type(basic const& other) const
   {
      return inherited::compare_same_type(other);
   }






   class Inner_Prod : public basic
   {
      GINAC_DECLARE_REGISTERED_CLASS(Inner_Prod, basic);

   public:
      ex lhs;
      ex rhs;

      Inner_Prod(ex l, ex r)
       : inherited(&Inner_Prod::tinfo_static)
       , lhs(l)
       , rhs(r)
      {
      }

      ex expand(unsigned options = 0) const
      {
         return expand_vector_product(this->lhs, this->rhs, inner_prod, options);
      }

      ex subs(exmap const&m, unsigned options = 0) const
      {
         ex subex = inherited::subs(m, options);
         if(!is_a<Inner_Prod>(subex))
            return subex;

         Inner_Prod sub_self = ex_to<Inner_Prod>(subex);
         if(    !is_a<lst>(sub_self.lhs)
             || !is_a<lst>(sub_self.rhs)
             || sub_self.lhs.nops() != sub_self.rhs.nops()
           )
            return subex;

         ex res = 0;
         for(const_iterator litr = sub_self.lhs.begin(),
                            ritr = sub_self.rhs.begin();
             litr != sub_self.lhs.end();
             ++litr, ++ritr)
             res += (*litr) * (*ritr);
         return res;
      }

      size_t nops() const
      {
         return 2;
      }

      ex op(size_t i) const
      {
         switch (i) {
         case 0:
             return this->lhs;
         case 1:
             return this->rhs;
         default:
             throw std::range_error("Inner_Prod::op(): no such operand");
         }
      }

      ex& let_op(size_t i)
      {
         ensure_if_modifiable();

         switch (i) {
         case 0:
             return this->lhs;
         case 1:
             return this->rhs;
         default:
             throw std::range_error("Inner_Prod::op(): no such operand");
         }
      }

      void do_print(print_context const& c, unsigned level = 0) const;
   };

   GINAC_IMPLEMENT_REGISTERED_CLASS_OPT(
      Inner_Prod, basic, print_func<print_context>(&Inner_Prod::do_print));

   Inner_Prod::Inner_Prod()
    : inherited(&Inner_Prod::tinfo_static)
   {
   }

   Inner_Prod::Inner_Prod(archive_node const& n, lst& sym_lst)
    : inherited(n, sym_lst)
   {
   }

   void Inner_Prod::archive(archive_node& n) const
   {
      inherited::archive(n);
   }

   ex Inner_Prod::unarchive(archive_node const& n, lst& sym_lst)
   {
      return (new Inner_Prod(n, sym_lst))->setflag(status_flags::dynallocated);
   }

   int Inner_Prod::compare_same_type(basic const& other) const
   {
      if(!is_a<Inner_Prod>(other))
         return 1;

      Inner_Prod const& o = static_cast<Inner_Prod const&>(other);
      return 
         (    (this->lhs == o.lhs && this->rhs == o.rhs)
           || (this->lhs == o.rhs && this->rhs == o.lhs))
         ? 0 : 1;
   }

   void Inner_Prod::do_print(print_context const& c, unsigned level) const
   {
      // print_context::s is a reference to an ostream
      c.s << "(" << this->lhs << " . " << this->rhs << ")";
   }

   inline ex inner_prod(ex const& l, ex const& r)
   {
      return Inner_Prod(l, r);
   }





   class Cross_Prod : public basic
   {
      GINAC_DECLARE_REGISTERED_CLASS(Cross_Prod, basic);

   public:
      ex lhs;
      ex rhs;

      Cross_Prod(ex l, ex r)
       : inherited(&Cross_Prod::tinfo_static)
       , lhs(l)
       , rhs(r)
      {
      }

      size_t nops() const
      {
         return 2;
      }

      ex expand(unsigned options = 0) const
      {
         return expand_vector_product(this->lhs, this->rhs, cross_prod, options);
      }

      ex subs(exmap const&m, unsigned options = 0) const
      {
         ex subex = inherited::subs(m, options);
         if(!is_a<Cross_Prod>(subex))
            return subex;

         Cross_Prod sub_self = ex_to<Cross_Prod>(subex);
         if(    !is_a<lst>(sub_self.lhs)
             || !is_a<lst>(sub_self.rhs)
             || sub_self.lhs.nops() != sub_self.rhs.nops()
           )
            return subex;

         if(sub_self.lhs.nops() == 2)
            return
                 sub_self.lhs.op(0) * sub_self.rhs.op(1)
               - sub_self.lhs.op(1) * sub_self.rhs.op(0)
            ;
         if(sub_self.lhs.nops() == 3)
         {
            lst res;
            for(int i = 0; i < 3; ++i)
               res.append(
                    sub_self.lhs.op((i+1)%3) * sub_self.rhs.op((i+2)%3)
                  - sub_self.lhs.op((i+2)%3) * sub_self.rhs.op((i+1)%3)
               );
            return res;
         }
         return subex;
      }

      ex op(size_t i) const
      {
         switch (i) {
         case 0:
             return this->lhs;
         case 1:
             return this->rhs;
         default:
             throw std::range_error("Cross_Prod::op(): no such operand");
         }
      }

      ex& let_op(size_t i)
      {
         ensure_if_modifiable();

         switch (i) {
         case 0:
             return this->lhs;
         case 1:
             return this->rhs;
         default:
             throw std::range_error("Cross_Prod::op(): no such operand");
         }
      }

      void do_print(print_context const& c, unsigned level = 0) const;
   };

   GINAC_IMPLEMENT_REGISTERED_CLASS_OPT(
      Cross_Prod, basic, print_func<print_context>(&Cross_Prod::do_print));

   Cross_Prod::Cross_Prod()
    : inherited(&Cross_Prod::tinfo_static)
   {
   }

   Cross_Prod::Cross_Prod(archive_node const& n, lst& sym_lst)
    : inherited(n, sym_lst)
   {
   }

   void Cross_Prod::archive(archive_node& n) const
   {
      inherited::archive(n);
   }

   ex Cross_Prod::unarchive(archive_node const& n, lst& sym_lst)
   {
      return (new Cross_Prod(n, sym_lst))->setflag(status_flags::dynallocated);
   }

   int Cross_Prod::compare_same_type(basic const& other) const
   {
      if(!is_a<Cross_Prod>(other))
         return 1;

      Cross_Prod const& o = static_cast<Cross_Prod const&>(other);
      return (this->lhs == o.lhs && this->rhs == o.rhs) ? 0 : 1;
   }

   void Cross_Prod::do_print(print_context const& c, unsigned level) const
   {
      // print_context::s is a reference to an ostream
      c.s << "(" << this->lhs << " x " << this->rhs << ")";
   }

   inline ex cross_prod(ex const& l, ex const& r)
   {
      if(is_a_vector(l) && is_a_vector(r) && l == r)
         return 0;
      return Cross_Prod(l, r);
   }
}

int main()
{
   using namespace GiNaC;

   vector v1("v1"), v2("v2"), v3("v3");
   lst l21(1,2), l22(2,-1);
   int l21_d_l22 = 0, l21_x_l22 = -5;
   lst l31(1,2,3), l32(-3,2,-1), l33(4,2,1);
   lst l31_x_l32(-8,-8,8);
   int l31_d_l32_x_l33 = -40;
   lst l31_x_l32_x_l33(-25,26,-9);

   // test indexing
   {
      assert(v1[1].subs(lst(v1 == l31)) == 2);
   }

   // test inner and cross product of vectors
   {
#define INNER1(a, b, av, bv)  inner_prod(a, b).subs(lst(a == av, b == bv))
#define CROSS1(a, b, av, bv)  cross_prod(a, b).subs(lst(a == av, b == bv))

      // purley symbolic expressions that are substituted at the end
      assert(cross_prod(v1, v1) == 0);
      assert(INNER1(v1, v2, l21, l22) == l21_d_l22);
      assert(CROSS1(v1, v2, l21, l22) == l21_x_l22);
      assert(CROSS1(v1, v2, l31, l32) == l31_x_l32);

#define INNER2(l, a, av)  inner_prod(l, a).subs(lst(a == av))
#define CROSS2(l, a, av)  cross_prod(l, a).subs(lst(a == av))

      // values and symbolic expressions are mixed
      assert(cross_prod(l21, l21) == 0);
      assert(INNER2(l21, v1, l22) == l21_d_l22);
      assert(CROSS2(l21, v1, l22) == l21_x_l22);
      assert(CROSS2(l31, v1, l32) == l31_x_l32);

#define TRIPLE1(a, b, c, av, bv, cv) \
   inner_prod(a, cross_prod(b, c)).subs(lst(a == av, b == bv, c == cv))
#define TRIPLE2(a, b, c, av, bv, cv) \
   cross_prod(a, cross_prod(b, c)).subs(lst(a == av, b == bv, c == cv))

      assert(TRIPLE1(v1, v2, v3, l31, l32, l31) == 0);
      assert(TRIPLE1(v1, v2, v3, l32, l31, l31) == 0);
      assert(TRIPLE1(v1, v2, v3, l31, l32, l33) == l31_d_l32_x_l33);
      assert(TRIPLE2(v1, v2, v3, l31, l32, l33) == l31_x_l32_x_l33);
   }

   using std::cout;
   using std::endl;

   // vector components
   cout << v1[2] << endl; // v1[2]

   // various products
   cout << inner_prod(v1, v2) << endl; // (v1 . v2)
   cout << cross_prod(v1, v2) << endl; // (v1 x v2)
   cout << inner_prod(v1, cross_prod(v2, v3)) << endl; // (v1 . (v2 x v3))
   cout << cross_prod(v1, cross_prod(v2, v3)) << endl; // (v1 x (v2 x v3))

   // identities (nothing fancy yet)
   cout << cross_prod(v1, v1) << endl; // 0

   // expand
   ex e1 = inner_prod(v1, v2 + v3);
   cout << e1 << endl; // (v1 . (v2 + v3)) 
   cout << expand(e1) << endl; // (v1 . v2) + (v1 . v3);

   // substitution (expand must be used before substitution)
   lst l1(1,2,3), l2(1,0,1), l3(1,1,0);
   cout << v1[2].subs(v1 == l1) << endl; // 3
   cout << expand(e1).subs(lst(v1 == l1, v2 == l2, v3 == l3)) << endl; // 7

   return 0;
}