algebra

# maps ket -> ket
# 3|x> => 3|x>
# |number: 7.2> => 7.2| >  # NB: the space in the ket label.
# 2|number: 3> => 6| >     # We can't use just |> because it is dropped all over the place!
# 8|number: text> => 0| >  # so the maths eqn: 3a + 7
# |3.7> => 3.7| >          # in my notation is 3|a> + 7| >
# 3|5> => 15| >
def category_number_to_number(one):         # find better name!
  one = one.ket()
  cat, value = extract_category_value(one.label)
  try:
    n = float(value)
  except:
    if cat == 'number':                     # not 100% want to keep these two lines
      return ket(" ",0)
    return one
  return ket(" ",one.value * n) 

def algebra_add(one,two):
  return one + two

def algebra_subtract(one,two):
  return delete3(one,two)

def algebra_mult(one,two,Abelian=True):
  one = superposition() + one  # hack so one and two are definitely sp, not ket
  two = superposition() + two

  result = superposition()
  for x in one.data:           # another hack! We really need to define an iterator for superpositions.
    x = category_number_to_number(x)  
    for y in two.data:
      y = category_number_to_number(y)
      print("x*y",x,"*",y)
      labels = [ L for L in x.label.split('*') + y.label.split('*') if L.strip() != '' ]
      if Abelian:  
        labels.sort()
      label = "*".join(labels)
      if label == '':         # we can't have ket("",value), since it will be dropped.
        label = " "
      result += ket(label,x.value * y.value)
  return result

def algebra_power(one,two,Abelian=True):
  one = superposition() + one
  two = category_number_to_number(two)
  try:
    n = int(two.value)
  except:
    return ket(" ",0)

  if n <= 0:
    return ket(" ",1)

  result = one
  for k in range(n - 1):
    result = algebra_mult(result,one,Abelian)
  return result

# implement basic algebra:
def algebra(one,operator,two,Abelian=True):
  op_label = operator if type(operator) == str else operator.the_label()
  null, op = extract_category_value(op_label)

  if op not in ['+','-','*','^']:
    return ket(" ",0)

  if op == '+':
    return algebra_add(one,two)            # Abelian option here too?
  elif op == '-':
    return algebra_subtract(one,two)       # ditto.
  elif op == '*':
    return algebra_mult(one,two,Abelian)
  elif op == '^':
    return algebra_power(one,two,Abelian)
  else:
    return ket(" ",0)

-- 3x * 5y
sa: algebra(3|x>,|*>,5|y>)
15.000|x*y>

-- (3x^2 + 7x + 13) * 5y
sa: algebra(3|x*x> + 7|x> + 13| >,|*>,5|y>)
15.000|x*x*y> + 35.000|x*y> + 65.000|y> 

-- (2x^2 + 11x + 19) * (11y + 13)
sa: algebra(2|x*x> + 11|x> + 19| >,|*>,11|y> + 13| >)
22.000|x*x*y> + 26.000|x*x> + 121.000|x*y> + 143.000|x> + 209.000|y> + 247.000| >

-- (a + b)^2
sa: algebra(|a> + |b>,|^>,|2>)
|a*a> + 2.000|a*b> + |b*b>

-- (a + b)^5
sa: algebra(|a> + |b>,|^>,|5>)
|a*a*a*a*a> + 5.000|a*a*a*a*b> + 10.000|a*a*a*b*b> + 10.000|a*a*b*b*b> + 5.000|a*b*b*b*b> + |b*b*b*b*b>

-- (a + b + c + d)^3
sa: algebra(|a> + |b> + |c> + |d>,|^>,|3>)
|a*a*a> + 3.000|a*a*b> + 3.000|a*a*c> + 3.000|a*a*d> + 3.000|a*b*b> + 6.000|a*b*c> + 6.000|a*b*d> + 3.000|a*c*c> + 6.000|a*c*d> + 3.000|a*d*d> + |b*b*b> + 3.000|b*b*c> + 3.000|b*b*d> + 3.000|b*c*c> + 6.000|b*c*d> + 3.000|b*d*d> + |c*c*c> + 3.000|c*c*d> + 3.000|c*d*d> + |d*d*d>

# simple complex number mult:
def complex_algebra_mult(one,two):
  one = superposition() + one  # hack so one and two are definitely sp, not ket
  two = superposition() + two

  result = superposition()
  for x in one.data:
    for y in two.data:
      if x.label == 'real' and y.label == 'real':
        result += ket("real",x.value * y.value)

      if x.label == 'real' and y.label == 'imag':
        result += ket("imag",x.value * y.value)

      if x.label == 'imag' and y.label == 'real':
        result += ket("imag",x.value * y.value)

      if x.label == 'imag' and y.label == 'imag':
        result += ket("real",-1 * x.value * y.value)
  return result

f |*> #=> 3 algebra(""|_self>,|^>,|2>) + 5 ""|_self> + 19| >

-- now, let's test it:
-- first a simple one:
sa: |x> => |a>
sa: f |x>
3.000|a*a> + 5.000|a> + 19.000| >

-- now a slightly more interesting one:
sa: |x> => |a> + |b>
sa: f |x>
3.000|a*a> + 6.000|a*b> + 3.000|b*b> + 5.000|a> + 5.000|b> + 19.000| >

-- define our functions:
sa: f |*> #=> 3 algebra(""|_self>,|^>,|2>) + 5 ""|_self> + 19| >
sa: g |*> #=> 7 algebra(""|_self>,|^>,|2>)

-- define x:
sa: |x> => |a>

-- f(x):
sa: f |x>
3.000|a*a> + 5.000|a> + 19.000| >

-- g(x):
sa: g |x>
7.000|a*a>

-- now a first try at g(f(x)):
sa: g f |x>
|>
-- Doh! I didn't expect that.
-- The explanation is the "" |_self> applied to |a*a>, |a> and then | >.

-- now, try again, this time indirectly:
sa: |y> => f |x>
sa: g |y>
63.000|a*a*a*a> + 210.000|a*a*a> + 973.000|a*a> + 1330.000|a> + 2527.000| >

-- NB: |_self> and not "" |_self>
sa: h |*> #=> 2 algebra(|_self>,|^>,|3>)
sa: h |x>
2.000|x*x*x>
sa: h (|a> + |b>)
2.000|a*a*a> + 2.000|b*b*b>

-- f(x) = 3x^2 + 5x +  19
sa: f |*> #=> 3 algebra(""|_self>,|^>,|2>) + 5 ""|_self> + 19| >
sa: |x> => |a>
sa: |y> => |a> + |b>
sa: |z> => |a> + |b> + |c>
sa: display-f |*> #=> display-algebra f |_self>
sa: table[ket,display-f] split |x y z>
+-----+----------------------------------------------------------------------+
| ket | display-f                                                            |
+-----+----------------------------------------------------------------------+
| x   | 3*a*a + 5*a + 19                                                     |
| y   | 3*a*a + 6*a*b + 3*b*b + 5*a + 5*b + 19                               |
| z   | 3*a*a + 6*a*b + 6*a*c + 3*b*b + 6*b*c + 3*c*c + 5*a + 5*b + 5*c + 19 |
+-----+----------------------------------------------------------------------+