generator.py

"""
 _____________________________________________
/ Roadmap 

Steps:
 # atom generation
    - probabilities of operations being used
 - atom simplification
 - atom filtering - no filtering needed! what are bad atoms?
 - atom combination
 - expression filtering
 - solution generation
 - equation filtering

            overall nesting level parameter?
            multiple combinations of same atoms?

            this is really just randomising structure itself

actual steps
 . expression type chosen (poly, trig, log)
 # atom generation & simplification
    . nesting modified based on expr type
 # atom combination
    . params modified based on expr type
    - intelligent combination
 . expression filtering
 . transforation if required (to trig or log)
 . solution generation
 . equation filtering
 
Doing:
 expression filtering
 solution generation
 equation filtering

parameters can be generated by user input + genetic algorithm

parameters:
    atom generation
    atom combination
    trigChance
    exponentialChance
    expression filter
    equation filter

e.g.
 - atom type probabilities,
 - expression type probabilities,
 - nesting level of problem,
 - problem difficulty (through depth of solution, depth of problem)

 _____________________________________________
/ Atom generation / combination 

Atom types? Trig atoms, exponential atoms?
Look through 3u logarithms and sin to see if there's much algebraic manipulation
Or should we just focus on identities and not put them randomly in equations
Should trig and exponential functions be in the same atom?

Should expressions have types as well?
polynomial, trig, exponential

Idea of atom combination + substitution
e.g. replace every x with e^x after expression generated
so just generate polynomials

How to deal with identities in this process?
testing stuff like: 
    a^2 - b^2 = (a-b)(a+b)
    2sin(x)cos(x) = cos(2x)
    etc.

Should we reduce scope to just polynomials?

 _____________________________________________
/ Types of atoms

 - polynomial (power, simple)
 - exponential (exponential)
 - trigonometry (trigonometry)

 exponentials and trigonometry -> substituted polynomials

"""

from sympy import *

inIPython = False
try:
    from IPython.display import display
    inIPython = True
except ImportError:
    display = pprint

import random

import Operations
from Timeout import exit_after

DEBUG = True

class mathGenerator:

    fillerNumRange = (1, 20)
    atomNesting = (1, 4)
    atomsInExpr = (1, 4)
    solutionRange = (1, 100)

    maxSolOps = 8
    acceptableBiggest = 10001


    trigChance = 0.1
    exponentialChance = 0.2

    def __init__(self, symbol=symbols('x')):
        self.symbol = symbol
        self.expr = symbol 

        # types are simple, trigonometry, exponential, power
        self.opps = Operations.allOperations

    def generate(self):
        numAtoms = random.randint(*mathGenerator.atomsInExpr)

        # generate atoms
        atoms = [self.genAtom(self.symbol) for _ in range(numAtoms)]

        # combine atoms
        expr = self.combineAtoms(atoms)

        # evaluate difficulty
        difficulty = self.evaluateExprDifficulty(expr)

        # generate solution
        solution = self.generateSolution(expr)

        # if unsolvable or some other error, try generation again
        if solution is None:
            if DEBUG:
                print('FATAL: failed to find solution, restarting')
            return self.generate()

        finalEquation, answer = solution

        return (finalEquation, answer, difficulty)

    @exit_after(1)
    def _solveInTermsOfA(self, expr):
        a, x = symbols('a x')
        sol = solve(Eq(expr, a), x)
        return sol

    @exit_after(1)
    def _findConstant(self, aExpr, answer):
        a = symbols('a')
        solutions = solve(Eq(aExpr, answer), a) 
        return solutions

    @exit_after(1)
    def _timedSolve(self, eq):
        return solve(eq, symbols('x'))

    def fuzzSolution(self, expr):
        # doesn't really fuzz, just tries to find roots
        
        failedToFindRoot = False

        sol = []
        tryRoots = Eq(expr, 0)
        try:
            sol = self._timedSolve(tryRoots)
            if not sol:
                if DEBUG:
                    print('RETRY: fuzz did not find values')
                failedToFindRoot = True
        except KeyboardInterrupt:
            if DEBUG:
                print('RETRY: fuzz took too long')
            failedToFindRoot = True
            pass

        # remove solutions with imaginary numbers 
        newSol = []
        for s in sol:
            if not (I in preorder_traversal(s)):
                newSol.append(s)
        sol = newSol

        if not sol:
            if DEBUG:
                print('RETRY: fuzz only found imaginary solutions')
            failedToFindRoot = True
            return None

        # remove solutions with too many operations
        noSolOps = [count_ops(s) for s in sol]
        if sol and max(noSolOps) > mathGenerator.maxSolOps:
            if DEBUG:
                print('RETRY: fuzz solutions too complicated')
            failedToFindRoot = True

        # remove solutions with massive numbers
        elif sol:
            biggestVal = max([max(self.leaves(s)) for s in sol])
            if biggestVal > mathGenerator.acceptableBiggest:
                if DEBUG:
                    print('RETRY: numbers in fuzz solution are too big')
                failedToFindRoot = True

        if failedToFindRoot:
            return None

        finalEquation = tryRoots
        solution = sol
        return (finalEquation, solution)

    def leaves(self, expr):
        if len(expr.args) == 0:
            return [expr]

        l = []
        for subtree in expr.args:
            l = l + self.leaves(subtree)
        return l

    def generateSolution(self, expr):
        a, x = symbols('a x')

        # attempt to find a general solution
        findingGeneralSolutionFailed = False

        try:
            sol = self._solveInTermsOfA(expr)
            if not sol:
                if DEBUG:
                    print('RETRY: could not find general solution')
                findingGeneralSolutionFailed = True
        except KeyboardInterrupt: # solution is taking too long to find
            findingGeneralSolutionFailed = True
            if DEBUG:
                print('RETRY: finding general solution took too long')
            pass


        # if can't find general solution (e.g. general quintic solution)
        # fuzz constant values and try to solve
        if findingGeneralSolutionFailed:
            return self.fuzzSolution(expr)

        # --- so at this point we have a general solution for a ---

        solutionInA = sol[0] # use first solution only to find a value

        # choose answer (for first solution)
        answer = random.randint(*mathGenerator.solutionRange)

        # attempt to find the right value of a to get chosen answer
        findingConstantValueFailed = False

        try:
            aVals = self._findConstant(solutionInA, answer)
            if not aVals:
                if DEBUG:
                    print('RETRY: could not solve for constant')
                findingConstantValueFailed = True
        except KeyboardInterrupt: # constant value is taking too long to find
            findingConstantValueFailed = True 

            if DEBUG:
                print('RETRY: finding constant took too long')
            pass

        # if have general solution but can't solve to specific answer,
        # fuzz again
        if findingConstantValueFailed:
            return self.fuzzSolution(expr)

        # -- at this point we have the constant value and an answer --

        # if multiple constant values end up with the same solution, choose one
        constantValue = random.choice(aVals)

        finalEquation = Eq(expr, constantValue)

        solutions = [answer]
        for otherSolution in sol[1:]:
            solutions.append(otherSolution.subs(a, constantValue))

        # -- validate all solutions --

        # remove solutions with imaginary numbers 
        newSol = []
        for s in solutions:
            if not (I in preorder_traversal(s)):
                newSol.append(s)
        solutions = newSol

        if not solutions:
            if DEBUG:
                print('RETRY: only found imaginary solutions')
            return None

        # remove solutions with too many operations
        noSolOps = [count_ops(s) for s in solutions]
        if solutions and max(noSolOps) > mathGenerator.maxSolOps:
            if DEBUG:
                print('RETRY: solutions too complicated')
            return None

        # remove solutions with massive numbers
        elif solutions:
            biggestVal = max([max(self.leaves(s + Integer(0))) for s in solutions])
            if biggestVal > mathGenerator.acceptableBiggest:
                if DEBUG:
                    print('RETRY: numbers in solution are too big')
                return None

        # finally return

        return (finalEquation, solutions)

    def evaluateExprDifficulty(self, expr):
        """
        difficulty = degree * no# operations
        """
        return degree(expr) * count_ops(expr)

    def combineAtoms(self, atoms):
        expr = atoms[0]
        for atom in atoms[1:]:
            funcName, funcInfo = self.chooseRandomOpp()

            # only choose a function that can combine atoms
            # also this is dumb (should be method to get atomCombiners and choose from them)
            while not funcInfo['atomCombiner']: # should the same opp be allowed twice in a row
                funcName, funcInfo = self.chooseRandomOpp()

            # randomise order of expr, atom applied to combiner
            if random.choice([True, False]):
                expr = funcInfo['func'](expr, atom) 
            else:
                expr = funcInfo['func'](atom, expr)

        return expr


    def genAtom(self, symbol=symbols('x')):
        # start with just the symbol
        atom = symbol

        # random level of nesting
        nest = random.randint(*mathGenerator.atomNesting)
        alreadyApplied = ['_']

        # returns list of complex function types applied
        complexFuncTypeAlreadyApplied = lambda t: t in \
                [info['type'] for _, info in alreadyApplied[1:] \
                if info['type'] != 'simple']

        for x in range(nest):
            # choose next function to apply
            funcName, funcInfo = self.chooseRandomOpp()

            # don't apply same function twice in a row
            # don't apply more than one of each trig, exponential or power function
            while alreadyApplied[-1][0] == funcName or complexFuncTypeAlreadyApplied(funcInfo['type']):
                funcName, funcInfo = self.chooseRandomOpp()

            alreadyApplied.append((funcName, funcInfo))

            # if double input function, generate the second input
            if funcInfo['numInputs'] == 2:
                # generate random number for second input
                secondVal = random.choice(range(*mathGenerator.fillerNumRange))

                # if double input function has a generator, use that instead
                if 'secondInputGen' in funcInfo.keys():
                    secondVal = (funcInfo['secondInputGen'])()

                atom = (funcInfo['func'])(atom, secondVal)

            # if single input expression, pass atom into it
            elif funcInfo['numInputs'] == 1:
                atom = (funcInfo['func'])(atom)

            # always simplify  #### WILL BE TIME HOG
            atom = simplify(atom)

        return atom

    # use with operationsOfType to choose random operation of type
    # return (name, info)
    def chooseRandomOpp(self, opps=Operations.allOperations):
        return random.choice(list(opps.items()))

if __name__ == "__main__":
    init_printing(use_unicode=True)

    generator = mathGenerator(symbols('x'))
    question, solution, difficulty = generator.generate()

    print()
    print("difficulty: " + str(difficulty))
    display(question)
    display(solution)