module Simulate where

import Data.Array
import Data.List
import Prelude hiding (null, length, or, foldr, maximum, concat, concatMap, foldl, foldr1, foldl1, sum, all, elem)
import Data.Ratio

import RandomFix
import LinearAlgebra
import Types

simulate :: Circuit -> Seed -> Time -> Time -> Temperature -> RBC -> [(Time, Output)]
simulate circuit seed dt end_time temp rbc	= select trace
  where
  	select					= takeWhile (\(t, _) -> t <= end_time) . map (\(t, s, o, r) -> (t, o))
	trace					= (concat . iterate step) ((tunnel . correct . solve . sources) (0, initial', initial', randoms))
	step					= tunnel . correct . solve . sources . integrate . last
	randoms					= map ((/ (2 ^ 30)) . fromInteger) (random (0, 2 ^ 30 - 1) seed)

	e					= 1.6021773e-19
	k					= 1.38066e-23
	dt'					= circa dt

	(initial', solve', integrate', correct_state', backdate')	= prepare circuit dt

	solve     (t, s, o, r) 			= (t, s, apply solve' s, r)
	integrate (t, s, o, r) 			= (t + dt, apply integrate' s, o, r)
	correct   (t, s, o, r)			= (t, apply correct_state' s, v_sub o (apply backdate' s), r)

	sources   (t, s, o, r)			= (t, random_offsets // map elements circuit, o, drop size r)
	  where
	 	elements (i, p, m, Conductor  _    )	= (i, 0)
		elements (i, p, m, Resistor   _    )	= (i, 0)
		elements (i, p, m, Capacitor  _    )	= (i, s ! i)
		elements (i, p, m, Inductor   _    )	= (i, s ! i)
		elements (i, p, m, Vsource    list )	= (i, source list)
		elements (i, p, m, Isource    list )	= (i, source list)
		elements (i, p, m, Junction   _    )	= (i, s ! i)

		random_offsets				= listArray (1, size) (map ((/ dt') . (* e). (* rbc) . (* 2) . (+ (-0.5))) r)
		(1, size)				= bounds s

		t'					= circa t

		source [(t0, v0)]					= v0
		source ((t0, v0) : _)			| t' < t0	= v0
		source ((t0, v0) : (t1, v1) : rest)	| t0 >= t1	= error "sources must have increasing times"
							| t' < t1	= v0 + (v1 - v0) * (t' - t0) / (t1 - t0)
							| otherwise	= source ((t1, v1) : rest)

	tunnel    (t, s, o, r)			= select (zip cumulative_probabilities new_states)
	  where
		e2					= e * e
		kT					= k * temp

		e0					= energy_state s
		energy_state  s				= sum (map (energy_state'  s  ) circuit)
		energy_source s o			= sum (map (energy_source' s o) circuit)

		events					= (concat . map events') circuit
		new_states				= map  (\(_,s,_,_) -> s) events
		delta_energies				= map ((\(r,s,o,_) -> (r, energy_state s - e0 + energy_source s o)) . correct) events
		probabilities				= limit_total (map (no_negative . probability) delta_energies)
		probability (r, de)	| r    == 0	= error "tunnel resistance must be nonzero (and should be >> 26k anyway)"
				 	| temp == 0	= dt' * (-de) / (e2 * circa r)
					| de   == 0	= dt' * kT / (e2 * circa r)
					| otherwise	= dt' * (-de) / (e2 * circa r * (1 - exp(de / kT)))
		no_negative p				= if p < 0 then 0 else p
		limit_total ps				= let total = sum ps in if total > 1 then map (/ total) ps else ps
		cumulative_probabilities		= tail (scanl (+) 0 probabilities)

		select []				= [(t, s, o, r')]
		select ((p', s') : zs)	| p < p'	= [(t, s, o, r'), (correct . solve) (t, s', o, r')]
					| otherwise	= select zs

		(p : r')				= r

		energy_state'  s   (i, p, m, Capacitor (c, _)   )	| c == 0	= error "capacitance must be nonzero"
									| otherwise	= 1/2 * (s ! i) ^ 2 / circa c
		energy_state'  s   (i, p, m, Junction  (c, _, _))	| c == 0	= error "junction capacitance must be nonzero"
									| otherwise	= 1/2 * (s ! i) ^ 2 / circa c
		energy_state'  _   _							= 0

		energy_source' s o (i, p, m, Vsource   _        )	= (s ! i) * (o ! i) * dt'
		energy_source' _ _ _					= 0

		events'        (i, p, m, Junction  (c, r, _))		= [(r, s // [(i, s ! i + e)], v_map (const 0) o, []),
									   (r, s // [(i, s ! i - e)], v_map (const 0) o, [])]
		events'        _					= []

prepare :: Circuit -> Time -> (Vector Approx, Matrix Approx, Matrix Approx, Matrix Approx, Matrix Approx)
prepare circuit dt = (v_map circa initial, m_map circa solve, m_map circa integrate, m_map circa correct_state, m_map circa backdate)
  where
	elements				= matrix element
	derivative				= matrix deriv
	initial					= vector init

	(l1, u1, i1)				= gauss_jordan               elements
	(l2, u2, i2)				= gauss_jordan (m_transpose  elements)
	constraint				=                            m_select_rows l1 i1
	freedom					=               m_transpose (m_select_rows l2 i2)

	(l3, u3, i3)				= gauss_jordan (m_mul constraint (m_mul derivative freedom))
	inverse3		| l3 == []	= i3
				| otherwise	= error "illegal circuit"
	correct			| l1 == []	= m_zero size
				| otherwise	= m_mul freedom (m_mul inverse3 constraint)
	correct_state				= m_sub (m_unit size) (m_mul derivative correct   )
	correct_output				= m_sub (m_unit size) (m_mul correct    derivative)

	solve					= m_mul correct_output (m_mul i1 correct_state)

	scaled_derivative			= m_map (* (dt / 2)) (m_mul derivative solve)
	(l4, u4, i4)				= gauss_jordan (m_sub (m_unit size) scaled_derivative)
	integrate		| l4 == []	= m_mul i4     (m_add (m_unit size) scaled_derivative)
				| otherwise	= error "unlucky choice of dt - try a smaller one"

	backdate				= m_map (/ dt) correct

	matrix f				= accumArray (+) (0 % 1) ((1, 1), (size, size)) ((filter no_ground . concat . map f) circuit)
	vector f				= accumArray (+) (0 % 1) ( 1    ,  size       ) ((                   concat . map f) circuit)

  	element	(i, p, m, Conductor  g       )	= [((p, i), -1 % 1), ((m, i),  1 % 1), ((i, m), -g    ), ((i, p),  g    ), ((i, i), -1 % 1)]
  	element	(i, p, m, Resistor   r       )	= [((p, i), -1 % 1), ((m, i),  1 % 1), ((i, m), -1 % 1), ((i, p),  1 % 1), ((i, i), -r    )]
	element	(i, p, m, Capacitor (c, _)   )	= [((p, i), -1 % 1), ((m, i),  1 % 1), ((i, m), -c    ), ((i, p),  c    )]
	element	(i, p, m, Inductor  (l, _)   )	= [((p, i), -1 % 1), ((m, i),  1 % 1), ((i, i),  l    )]
	element	(i, p, m, Vsource    _       )	= [((p, i), -1 % 1), ((m, i),  1 % 1), ((i, m), -1 % 1), ((i, p),  1 % 1)]
	element	(i, p, m, Isource    _       )	= [((p, i), -1 % 1), ((m, i),  1 % 1), ((i, i),  1 % 1)]
	element	(i, p, m, Junction  (c, _, _))	= [((p, i), -1 % 1), ((m, i),  1 % 1), ((i, m), -c    ), ((i, p),  c    )]

  	deriv 	(i, p, m, Conductor  _       )	= []
  	deriv 	(i, p, m, Resistor   _       )	= []
	deriv 	(i, p, m, Capacitor (_, _)   )	= [((i, i),  1)]
	deriv 	(i, p, m, Inductor  (_, _)   )	= [((i, p),  1), ((i, m), -1)]
	deriv 	(i, p, m, Vsource    _       )	= []
	deriv 	(i, p, m, Isource    _       )	= []
	deriv 	(i, p, m, Junction  (_, _, _))	= [((i, i),  1)]

  	init	(i, p, m, Conductor  _       )	= []
  	init   	(i, p, m, Resistor   _       )	= []
	init    (i, p, m, Capacitor (_, q)   )	= [(i, q)]
	init   	(i, p, m, Inductor  (_, f)   )	= [(i, f)]
	init   	(i, p, m, Vsource    _       )	= []
	init   	(i, p, m, Isource    _       )	= []
	init   	(i, p, m, Junction  (_, _, q))	= [(i, q)]

	no_ground ((r, c), _)			= r /= 0 && c /= 0
	size					= maximum (concat (map (\(p, m, i, _) -> [p, m, i]) circuit))
