# use that in Jupiter web client
def mdisplay_md(fmt, *args, **kwargs):
    display(Markdown(fmt.format(
        *[f'${latex(x)}$' for x in args],
        **{k: f'${latex(v)}$' for k, v in kwargs.items()})
    ))

def as_text(thing):
    from IPython.core.interactiveshell import InteractiveShell # type: ignore
    plain_formatter = InteractiveShell.instance().display_formatter.formatters['text/plain']
    pp = plain_formatter(thing)
    lines = pp.splitlines()
    return lines

def vpad(lines, height):
    width = len(lines[0])
    return [' ' * width for _ in range(height - len(lines))] + lines

# terminal and emacs can't display markdown, so we have to use that as a workaround
def mdisplay_plain(fmt, *args, **kwargs):
    import re
    from itertools import chain
    fargs   = [as_text(a) for a in args]
    fkwargs = {k: as_text(v) for k, v in kwargs.items()}

    height = max(len(x) for x in chain(fargs, fkwargs.values()))

    pargs   = [vpad(a, height) for a in fargs]
    pkwargs = {k: vpad(v, height) for k, v in fkwargs.items()}

    pos = height // 2

    lines = []
    for h in range(height):
        largs   = [a[h] for a in pargs]
        lkwargs = {k: v[h] for k, v in pkwargs.items()}
        fstr = fmt if h == pos else re.sub('[^{}]', ' ', fmt)
        lines.append(fstr.format(*largs, **lkwargs))
    print('\n'.join(lines))

mdisplay = mdisplay_plain

%matplotlib inline
# TODO P hmm try that? %matplotlib notebook . that fucks up export though...
from IPython.display import display, Math, Markdown
from sympy import *
init_printing(use_latex='mathjax')
# TODO P wonder if we can convert to latex by default??
# TODO P I guess the answer is here?? https://docs.sympy.org/latest/_modules/sympy/printing/str.html
# TODO P shit. text has to be displayed as markdown...
# TODO P shit. maybe wrap patterns in text, and everythin else is latex??
# TODO render line numbers
# somehow there is no nice function in sympy to display stuff in the same line...

from sympy import Derivative as D, Function as F
x, t = symbols('x t', real=True)
u = F('u')(x, t)
eq = Eq(D(u, x, 2), D(u, t, 2))
x_l = 0
L = symbols('L', real=True)
x_r = L

mdisplay("PDE: {}", eq) # TODO velocity??
mdisplay("BCs: $u(x_l, t) = 0$, $u(x_r, t) = 0$, $x_l=${}, $x_r=${}", x_l, x_r)

# pdsolve(eq)
# TODO shit! sympy stops at first error...
# TODO find a way to ignore the errors

display(pde.allhints)

X = F('X')(x)
T = F('T')(t)
[sep_X, sep_T] = pde_separate_mul(eq, u, [X, T])
display((sep_X, sep_T))
# TODO at which point are we actually invoking linearity??

E = symbols('E', real=True)
eqX = Eq(sep_X, -E)
eqT = Eq(sep_T, -E)
display((eqX, eqT))

classify_ode(eqX)

bc_l = Eq(X.subs(x, x_l), 0)
bc_r = Eq(X.subs(x, x_r), 0)
mdisplay("BCS: {}, {}", bc_l, bc_r)
Xsol = dsolve(
    eqX,
    X,
    ics={
        bc_l.lhs: bc_l.rhs,
        bc_r.lhs: bc_l.rhs,
    },
    hint='nth_linear_constant_coeff_homogeneous',
)
display(Xsol)

Xsol = dsolve(
    eqX,
    X,
    hint='nth_linear_constant_coeff_homogeneous',
)
display(Xsol)

constants = list(Xsol.free_symbols - eqX.free_symbols)
system = [
    Eq(Xsol.subs(x, x_l).rhs, 0),
    Eq(Xsol.subs(x, x_r).rhs, 0),
]
display(system)

linsolve(system, constants)

from sympy.solvers.solveset import linsolve as slinsolve
slinsolve(system, constants)

from sympy.solvers.solveset import nonlinsolve as snonlinsolve
nonlinsolve(system, constants + [E])

Y = symbols('Y')
system = [
    Y * (exp(sqrt(-E)) - 1)
]
display(*system)

snonlinsolve(system, [Y, E])

solveset(exp(sqrt(-E)) - 1)
# TODO need to store captured outputs for persistence

Y = symbols('Y')
solveset(exp(Y) - 1)

from sympy.solvers.solveset import solveset_complex
solveset_complex(sqrt(1) + Y, Y)

k = symbols('k')
subst = Eq(sqrt(-E), k)

system = [
    eq.subs(subst.lhs, subst.rhs) for eq in system
]
Xsol = Xsol.subs(subst.lhs, subst.rhs)
display(system)
display(Xsol)

res = snonlinsolve(system, constants + [k])
mdisplay("Solutions: {}", res)

def bindFiniteSet(x):
    assert isinstance(x, FiniteSet)
    return x.args
def bindImageSet(x):
    assert isinstance(x, ImageSet)
    return x.lamda, x.base_set
(first, second) = bindFiniteSet(res)
(C1sol, C2sol, kSol) = second
lamda, base = bindImageSet(kSol)
mdisplay("$C_1$ = {}; C2 = {}; k = {}", C1sol, C2sol, kSol)

def bindTuple(x):
    assert isinstance(x, Tuple)
    return x.args

def flatten(x):
    args = bindTuple(x)
    args = [x if isinstance(x, Set) else FiniteSet(x) for x in args]
    return ProductSet(*args)

def collapse(solutions):
    ret = EmptySet()
    for sol in solutions.args:
        ret = ret.union(flatten(sol))
    return ret

collapsed = collapse(res)
display(collapsed)

# TODO need to turn to ImageSet?
# TODO and then we can substitute for each N separately...

Xsol = Xsol.subs({
    constants[0]: C1sol,
    constants[1]: C2sol,
})
# only one constant 
# [const] = set(constants) - Xsol.free_symbols
# display(const)
A = IndexedBase('A')
B = IndexedBase('B')
n = symbols('n', integer=True)
term = Xsol.rhs.subs({
    constants[0]: A[n],
    constants[1]: B[n],
    k: lamda(n),
})
term = collect(term, [A[n], B[n]]) # for nicer grouping
basis = ImageSet(
    Lambda(n, term),    
    base,
)
mdisplay("Solutions basis: {}", basis) # TODO shit, that doesn't need to be formatted like latex

def eliminate_dependency(iset, idx, var):
    assert iset.base_set == Integers
    prop = simplify(basis.lamda(idx) / basis.lamda(-idx))
    print(prop.free_symbols)
    assert var not in prop.free_symbols
    # now that we verified linear dependency, we can shrink the index space
    return ImageSet(iset.lamda, Naturals0)

# basis = eliminate_dependency(basis, n, x)
# display(basis)

# shit, that didn't work. requires lower and upper bounds
# kinda makes sense, we can't just sum over arbitrary sets...
# but we do want to sum over countable sets!
def sum_iset(iset):
    bs = iset.base_set
    lamda = iset.lamda
    [var] = lamda.variables
    return Sum(
        lamda(var),
        (var, bs._inf, bs._sup),
    )
Xsol2 = sum_iset(basis)
display(Xsol2)

# TODO FIXME exploiting that A = 0..
def get_func(idx):
    return simplify(Xsol2.function.subs(n,idx) / B[idx])

# if we not assume n to be nonnegative it will make integration a bit easier for sympy
# from sympy.assumptions.assume import global_assumptions
# from sympy import Q
# global_assumptions.add(Q.nonnegative(n))
m = symbols('m', integer=True)# , nonnegative=True)

# TODO ???? takes a lot to compute this

# integrate(get_func(n) * get_func(m), (x, 0, L))
# TODO wtf??? that's messy

#def sum_map(f, iset, idx):
#    return Sum(
#        f(idx, iset.lamda(idx)),
#        (idx, iset.base_set._inf, iset.base_set._sup),
#    )


# B = IndexedBase('B')

# sol = sum_map(lambda idx, term: Xsol.subs({
#    C1: A[idx],
#    C2: B[idx],
#    k : term,
#}), K, idx=n)


# print('General solution:')
# display(sol)

from sympy.solvers.solveset import invert_complex, invert_real, _solve_radical, _solveset, solveset_complex
QQ = symbols('QQ', real=True)
# snonlinsolve([(Y - 1) * (exp(I * sqrt(YY)) - 1)],[Y, YY])
# TODO FIX -- simplify its output??
from sympy.solvers.solveset import linsolve as slinsolve
display(linsolve(system, constants))
# TODO shit is E even the right name for the constant???

from sympy.solvers.solveset import linsolve
from sympy.solvers.ode import ode_nth_linear_constant_coeff_homogeneous as ode_solve
from sympy.core.numbers import IntegerConstant
# `
k = symbols('k')
E = k ** 2
# E = symbols('E', real=True)
# eqT = Eq(_eqT + E) # ).canonical
# canonical moves all variables to the left and numbers to the right
Tsol = dsolve(eqT).rhs
# shit, it won't eliminate solutions based on k :(
Xsol = dsolve(
    eqX,
    X,
    #ics={
    #    X.subs(x, x_l): 0,
    #    X.subs(x, x_r): 0,            
    #},
    hint='nth_linear_constant_coeff_homogeneous',
    # returns = 'both',
    
).rhs

print("General solutions for X and T:")
display((Xsol, Tsol))

# how to specify constants for DE solver??
[C1, C2] = list(Xsol.free_symbols.difference({k, x}))

bcs_l = Eq(Xsol.subs(x, x_l))
bcs_r = Eq(Xsol.subs(x, x_r))

print("BC enforces:")
display((bcs_l, bcs_r))


# TODO blog about that, give examples?
# so, solveset stuff can't handle 
[c1_c2] = solve(bcs_l, C1)
# solve(bcs_r.subs(C1, c1_c2))
# [{C2:0},{k:0},{k:iπ}]
# TODO what do we do with C2 = 0 solution???
K = solveset(bcs_r.subs(C1, c1_c2), k)

print("Allowed values of k:")
display(K)

# TODO how to eliminate negative k solutions?
# there must be something wrong with our approach if we have to do that manually. right?
# something unnatural. maybe we want to impose some sort of symmetry first??
# ah, basically separate equation in left and right waves first?

# so, can we eliminate extra constants and instead use the whole range of ks ????

# hmm. maybe the redundancy is coming from separation of variables? e.g. negative and positive energies?


# TODO generate sequence of edits from git commits?...

# basis = imageset(lambda kk: Xsol.subs(k, kk), K)
# print("Solutions basis:")
# display(basis)

# here, sympy retuns something called ImageSet. it basically means that the solutions image of Z under x -> x i pi


# right, how to define infinite coefficients? can it just be a function?
# TODO how to sum something over imageset????

# TODO base_set??


# shit, that didn't work. requires lower and upper bounds
# kinda makes sense, we can't just sum over arbitrary sets...
# but we do want to sum over countable sets!

def sum_map(f, iset, idx):
    return Sum(
        f(idx, iset.lamda(idx)),
        (idx, iset.base_set._inf, iset.base_set._sup),
    )

A = IndexedBase('A')
B = IndexedBase('B')
n = symbols('n')

sol = sum_map(lambda idx, term: Xsol.subs({
    C1: A[idx],
    C2: B[idx],
    k : term,
}), K, idx=n)


print('General solution:')
display(sol)

# hmm, it's a bit awkward, we've got quite a bit of redundancy in 
# in textbooks, that's violated in an ad-hoc manner at the step (TODO refer to line, with a marker?)
# the proper way to deal with this would be rewriting the solution
# and only then enforce a bc!

# TODO do we need to eliminate negative coefficients??

# imageset(lambda u: u**2, allowed_ks)
#Xsol
# print(allowed_ks.domain)
# ok, now we basically restrict ourselves to considering only specific waves

# TODO pick first N coeffs and do initial transform?
# usol = Xsol.rhs * Tsol.rhs
# linsolve(, [C1, C2])
print("rewrite: C[0] = (A_0 + B_0) / 2; for n>0: C[n] = B_n + A_{-n}; C[-n] = B_{-n} + A_n")
# C = IndexedBase('C')
C = Function('C')
# sol = Sum(C[n] * exp(I * pi * n * x), (n, -oo, oo))

# TODO here, extract the exponent from split solution..
sol = Integral(C(k) * exp(I * k * x), (k, -oo, oo))
print("General sol")
display(sol)
# TODO how to check they are equivalent??
bc_l = sol.subs(x, x_l)
bc_r = sol.subs(x, x_r)
print("BCS:")
display((Eq(bc_l), Eq(bc_r)))
# TODO is there a way to solve that automatically???

def mComm(expr, cls):
    assert isinstance(expr, cls)
    return expr.args

def mAdd(expr):
    # TODO is there anything standard in sympy?
    return mComm(expr, Add)
    
def mMul(expr):
    return mComm(expr, Mul)

def mIntegral(expr):
    assert isinstance(expr, Integral), type(i)
    return expr.limits, expr.function

def pullIntegral(expr):
    if isinstance(expr, Integral):
        return expr
    if isinstance(expr, Mul):
        # TODO might work for a constant...
        parts = [pullIntegral(p) for p in mMul(expr)]
        # TODO only one should be integral, rest constants?
        consts = []
        rest = []
        for p in parts:
            if isinstance(p, IntegerConstant):
                consts.append(p)
            else:
                rest.append(p)
        rest = [pullIntegral(r) for r in rest]
        res = None
        if len(rest) == 1:
            lims, fun = mIntegral(rest[0])
            res = Integral(Mul(*consts) * fun, lims)
            return res
        else:
            res = Mul(*consts)
            if len(rest) > 0:
                res = res * Mul(*rest)
            return res
    return expr
    

def asIntegral(expr):
    res = pullIntegral(expr)
    assert isinstance(res, Integral), type(res)
    return res
    # if is

def combine_isum(isum):
    parts = mAdd(isum)
    parts = [asIntegral(p) for p in parts]
    # TODO shit, is there something for pulling stuff under integral?
    # TODO expression with limits
    lims = [i.limits for i in parts]
    exprs = [i.function for i in parts]
    [lim] = list(set(lims))
    return Integral(Add(*exprs), lim)
    
    # TODO generify for expr with limits
    
def int2sum(thing):
    lims, func = mIntegral(thing)
    [q1, _] = mMul(func)
    KK = solveset(q1, k)
    return KK
    
    # TODO how to bind it to list of integrals? some sort of generic shape??
# TODO right! so if we got independent things, we should test BOTH linear combinations
bcs1 = rcollect(combine_isum(bc_l + bc_r), C(k))
display(bcs1)

bcs2 = rcollect(combine_isum(bc_l - bc_r), C(k))
display(bcs2)

#lims1, func1 = mIntegral(bcs1)
#[qq1, _] = mMul(func1)
#lims2, func2 = mIntegral(bcs2)
#[qq2, _] = mMul(func2)

#from sympy.solvers.solveset import linsolve as slinsolve, nonlinsolve as snonlinsolve
# KK = solveset([qq1, qq2], k)

# KK = slinsolve([qq1, qq2], k)
# display(qq1)
# display(qq2)
# nope.. they were independent and should be solved independently!!

def the(objs):
    assert len(objs) > 0, "expected non-empty list"
    assert all([o == objs[0] for o in objs]), "not all objects are equal"
    return objs[0]

def asLambda1(expr):
    assert isinstance(expr, Lambda)
    return expr.variables, expr.expr

# TODO https://docs.sympy.org/latest/tutorial/manipulation.html is not too helpful though...
def combine_imagesets(s1, s2):
    display(s1)
    display(s2)
    sets = [s1, s2]
    # TODO assert all are image sets? or pull them out
    dom = the([s.base_set for s in sets])
    # TODO ok, need to  analyse lamdas...
    # TODO should be same variable???
    v1, e1 = asLambda1(s1.lamda)
    v2, e2 = asLambda1(s2.lamda)
    # print(cse([e1, e2]))
    # assert v1 == v2
    # basically, we need to figure out if e1 == e2 for some n? then we can discard one of them!!
    # and if only one elem left, we don't need the finiteset?
    # ah shit, we need to replace the function!
    # display(solve(e1 - e2))
    print("GCD")
    display(factor_list(e1))
    display(factor_list(e2))
    display(gcd_terms([e1, e2]))
    return imageset(lambda q: FiniteSet(*[s.lamda(q) for s in sets]), dom)
    
    
# TODO shit, how to simplify them???
# combine_imagesets()
# fucking hell. why is that so tricky????
# anyway:
display(int2sum(bcs) + int2sum(bcs2))
print("which is same as:")
# TODO again, need some sort of quickcheck for expressions, sets etc
n = symbols('n')
Ns = imageset(Lambda(n, pi * n), Integers)
display(Ns)
# TODO now we actually need int2sum
# something that pick subintegral expression 

def change_limits(expr, ns):
    # expr.limits
    ef = expr.function
    [ev] = expr.variables
    lamda = ns.lamda
    lf = lamda.expr
    [lv] = lamda.variables
    bs = ns.base_set   
    return Sum(ef.subs(ev, lamda(lv)), (lv, bs._inf, bs._sup))
# ugh. so dirty..
res2 = change_limits(sol, Ns)
print("After substituting restrictions on wavevectors")
display(res2)



from sympy.integrals.transforms import _fourier_transform

u = symbols('u')
integrate(res2 * exp(-I * pi * u * x), (x, x_l, x_r))

# shit. I guess we run at the fact that expr(-I pi u x) are not basis functions for the problem once again...
# still, shouldn't prevent us from solving. right???


# fucking hell.. why does that work???


# TODO next: initial condition to find Cs?
# basically, here we're going to restrict ourselves? or we can get away with fully symbolic for a while?...
# TODO ok, just multiply all by exp(- i pi k x)
# TODO how to eliminate C_{k+1} vs C_{k-1} ???
# wonder if division by two in classic particle in a box coming from the fact that they have to normalise by 2
# since they are using degenerate solution space


# TODO do final chage of variables from C(pi * n) to some indexed function?..

# TODO shit... but coeffs are not independent.. because originally we had two equations??
# mm. ok, individual things don't have to satisfy bcs!


    
# when k = 2 pi n, it works for any C
# when it's not, thet that can only work when all of C are 0, due to properties of L^2 space
# TODO could implement solving integral equations in python?...
# basically, C(k) is non-zero on k = 2 pi n and 0 otherwise

# TODO hmm interesting, looks kinda like a discrepancy between pi *n and 2 n pi + pi
# but, actually, I guess this is kinda disguised in the linear combination thing! this form explicitly forbids cosines!

# shit. it can't handle restrictions on k :(
from sympy.solvers.solveset import linsolve as slinsolve

A, B, k = symbols('A B k')


slinsolve([
    A * exp(sqrt(-k)) + exp(sqrt(k)),
], [A])

# right, it just can't figure out multiple solutions??
# solveset(exp(k) - 1)

# shit, it can't handle eliminating certain k's :( 
slinsolve([
    A,
    0,
    #exp(k) - 1,
], [A, B])

E = symbols('E', real=True)
k = symbols('k')

# TODO shit, it won't handle sqrt??
solveset(exp(k) - exp(-k))


# solveset(exp(k) - 1)
solveset(exp(sqrt(k**2)) - 1)

slinsolve([
    A + B,
    A * exp(- k) + B * exp(k),
],
    [A, B]
)


#solve([
#    exp(k) + 1,
#], returns='solveset')

from sympy.core.rules import Transform
def split(integ):
    return Add(*[integ.func(term, *integ.args[1:]) for term in Add.make_args(integ.args[0])])
F = Function('F')
G = Function('G')
H = Function('H')
A, B = symbols('A B')
z = symbols('z')
v = symbols('v')
it = Integral(F(z) * z + -G(z) * z, (z, 0, oo))
display(it)
thing = it.xreplace(Transform(split, lambda i: isinstance(i, Integral)))

[left, right] = mAdd(thing)
rrr = Add(
    left,
    right.transform(-z, z),
)
display(rrr)
rrr.subs_expr({
    F(z): H(z),
    G(-z): H(z),
})
# TODO shit...
# TODO replace function with a partial??



# TODO how to assert that expressions are equivalent???

q, w = symbols('q w')
ee = Eq(q, w)
ee.subs(q, 0)
# display(ee)

# TODO NOTE: the bc equations:
# solve [C1 + C2 =0, C1 * exp(i * k) + C2 * exp(-i * k) =0]
# sage couldn't hande it either
# wolfram alpha did though!