from typing import TYPE_CHECKING
from scipy import imag, real, log, sqrt, random, zeros
from scipy.interpolate import RectBivariateSpline
if TYPE_CHECKING:
    from QuantumSketchBook.schroedinger import Schroedinger


class Nelson:
    def __init__(self, schroedinger: "Schroedinger", n: int, micro_steps=10):
        self.mesh = schroedinger.mesh
        psi = schroedinger.solution()
        imz, rez = imag(log(psi)), real(log(psi))
        self.realSpline = RectBivariateSpline(self.mesh.t_vector, self.mesh.x_vector, rez)
        self.imagSpline = RectBivariateSpline(self.mesh.t_vector, self.mesh.x_vector, imz)
        self.x = schroedinger.x0state.random_values(n)
        self.t = self.mesh.t_min
        self.n = n
        self.t_micro = self.mesh.dt / micro_steps
        self.micro_steps = micro_steps

    def drift(self, x, t):
        re = self.realSpline.ev(t, x, dy=1)
        im = self.imagSpline.ev(t, x, dy=1)
        return re + im

    def run(self):
        rand = sqrt(0.5) * random.normal(scale=sqrt(self.t_micro), size=self.n)
        self.x += self.drift(self.x, self.t) * self.t_micro + rand
        self.t += self.t_micro
        return self.x

    def set_t(self, new_t):
        self.t = new_t
        return self

    def locus(self):
        locus_array = zeros((self.n, self.mesh.t_num))
        locus_array[:, 0] = self.x
        for i in range(1, self.mesh.t_num):
            print("\rNelson have solved {:.2%}!".format(i / self.mesh.t_num), end=" ", flush=True)
            for j in range(self.micro_steps-1):
                self.run()
            locus_array[:, i] = self.run()
        print("\rNelson have solved 100.00%!", flush=True)
        return locus_array

@@ -111,29 +111,21 @@
  import 
-RectBivariate
+Cubic
 Spline%0Ai
@@ -312,15 +312,8 @@
 elf.
-mesh = 
 schr
@@ -324,283 +324,61 @@
 nger
-.mesh%0A        psi = schroedinger.solution()%0A        imz, rez = imag(log(psi)), real(log(psi))%0A        self.realSpline = RectBivariateSpline(self.mesh.t_vector, self.mesh.x_vector, rez)%0A        self.imagSpline = RectBivariateSpline(self.mesh.t_vector, self.mesh.x_vector, imz)
+ = schroedinger%0A        self.mesh = schroedinger.mesh
 %0A   
@@ -433,41 +433,8 @@
 (n)%0A
-        self.t = self.mesh.t_min%0A
     
@@ -472,24 +472,37 @@
 icro = self.
+schroedinger.
 mesh.dt / mi
@@ -564,136 +564,169 @@
 ef d
+e
 ri
-ft(self, x, t):%0A        re = self.realSpline.ev(t, x, dy=1)%0A        im = self.imagSpline.ev(t, x, dy=1)%0A        return re + im
+vative_fo_part(self, part, sch):%0A        log_part = part(log(sch))%0A        spline = CubicSpline(self.mesh.x_vector, log_part)%0A        return spline(self.x, 1)
 %0A%0A  
@@ -735,27 +735,80 @@
 def run(self
+, drift
 ):%0A
+        for _ in range(self.micro_steps):%0A    
         rand
@@ -879,24 +879,28 @@
 .n)%0A        
+    
 self.x += se
@@ -897,21 +897,16 @@
 lf.x += 
-self.
 drift(se
@@ -913,16 +913,8 @@
 lf.x
-, self.t
 ) * 
@@ -945,113 +945,8 @@
     
-self.t += self.t_micro%0A        return self.x%0A%0A    def set_t(self, new_t):%0A        self.t = new_t%0A        
 retu
@@ -952,16 +952,18 @@
 urn self
+.x
 %0A%0A    de
@@ -1115,107 +1115,8 @@
 m):%0A
-            print(%22%5CrNelson have solved %7B:.2%25%7D!%22.format(i / self.mesh.t_num), end=%22 %22, flush=True)%0A
     
@@ -1315,8 +1315,394 @@
 s_array%0A
+%0A    def __iter__(self):%0A        self.x = self.schroedinger.x0state.random_values(self.n)%0A        yield self.x%0A        for i, sch in enumerate(self.schroedinger):%0A            drift = self.derivative_fo_part(real, sch) + self.derivative_fo_part(imag, sch)%0A            yield self.run(drift)%0A            print(f%22%5CrNelson have solved %7B(i + 1) / self.mesh.t_num:.2%25%7D!%22, end=%22 %22, flush=True)%0A

make nelson use less memory