module Common = struct
	let id a = a
	let dup a = a, a
	let some v = Some v
	
	let curry2 f a b = f (a, b)
	let uncurry2 f (a, b) = f a b
	let curry3 f = curry2 (curry2 f)
	let uncurry3 f = uncurry2 (uncurry2 f)
	
	let first a b = a
	let second a b = b
	
	let ( =* ) a b =
		try a = b with Invalid_argument _ -> a == b
	
	let ( <>* ) a b =
		try a <> b with Invalid_argument _ -> a != b
end

open Common

(* TODO: add necessary functions to Event (such as value change -> event) so
   that we may rewrite FrCore without manipulating events directly *)

module Event = struct
	type 'a t = 'a option
	
	let none = None

	let tag v = function
		| Some _ -> Some v
		| None -> None

	let attach v = function
		| Some a -> Some (v, a)
		| None -> None

	let map f = function
		| Some a -> Some (f a)
		| None -> None
	
	let may = function
		| Some a -> a
		| None -> None
	
	let lmerge a b = match a, b with
		| Some _, _ -> a
		| _, _ -> b

	let rmerge a b = match a, b with
		| _, Some _ -> b
		| _, _ -> a

	let merge f a b = match a, b with
		| Some a, Some b -> Some (f a b)
		| Some _, None -> a
		| None, Some _ -> b
		| None, None -> None

	let map_merge fa fb fab a b = match a, b with
		| Some a, Some b -> Some (fab a b)
		| Some a, None -> Some (fa a)
		| None, Some b -> Some (fb b)
		| None, None -> None

	let cat l =
		match (List.fold_left (fun r -> function
			| Some a -> a :: r
			| None -> r) [] l)
		with
		| [] -> None
		| r -> Some (List.rev r)

	let join a b = match a, b with
		| Some a, Some b -> Some (a, b)
		| _, _ -> None

	let split = function
		| Some (a, b) -> Some a, Some b
		| None -> None, None

	let filter f = function
		| Some a when f a -> Some a
		| _ -> None
	
	let map_filter f = function
		| Some a -> f a
		| None -> None
	
	let get e = function
		| Some _ -> Some e
		| None -> None

	(* internal *)
	let return a = Some a
	let map_default d f = function
		| Some a -> f a
		| None -> d
	let is_event = function Some _ -> true | None -> false
	let from_event = function Some a -> a | None -> assert false
end


type ('i, 'o) sf =
	| TF of ('i -> ('i, 'o) transition)
	| Const of 'o
and ('i, 'o) transition = ('i, 'o) sf option * 'o

exception Undefined_signal

let transition sf i =
	match sf with
	| TF tf -> begin
		match tf i with
		| Some sf', o -> sf', o
		| None, o -> sf, o
		end
	| Const o -> sf, o

let make_sf tf = TF tf

let is_constant = function Const _ -> true | _ -> false


(* Basic signals *)

let constant a = Const a

let arr f = make_sf (fun i -> None, f i)

let arr2 f = arr (uncurry2 f)

let rec (>>>) a b =
	make_sf (fun i ->
		let a', ao = transition a i in
		let b', o = transition b ao in
		match a', b' with
		| _, Const bo' ->
			Some (constant bo'), o
		| Const ao', _ ->
			Some (constant (snd (transition b' ao'))), o
		| _ ->
			Some (a' >>> b'), o)

let (<<<) a b = b >>> a

let rec first a =
	make_sf (fun (ai, bi) ->
			let a', ao = transition a ai in
			Some (first a'), (ao, bi))

let rec second b =
	make_sf (fun (ai, bi) ->
			let b', bo = transition b bi in
			Some (second b'), (ai, bo))

let rec ( *** ) a b =
	make_sf (fun (ai, bi) ->
			let a', ao = transition a ai
			and b', bo = transition b bi in
			match a', b' with
			| Const ao', Const bo' ->
				Some (constant (ao', bo')), (ao, bo)
			| _ ->
				Some (a' *** b'), (ao, bo))

let rec (&&&) a b =
	make_sf (fun i ->
			let a', ao = transition a i
			and b', bo = transition b i in
			match a', b' with
			| Const ao', Const bo' ->
				Some (constant (ao', bo')), (ao, bo)
			| _ ->
				Some (a' &&& b'), (ao, bo))

let returnA = TF (fun i -> None, i)

(* I doubt this works at all *)
(*let rec loop a =
	{
		tf = (fun ai ->
			let rec a'_ao_bo =
				lazy (transition a (ai, (lazy (snd (snd (Lazy.force a'_ao_bo)))))) in
			let a', (ao, bo) = Lazy.force a'_ao_bo in
			loop a', ao)
	}*)

let rec loop init a =
	make_sf (fun ai ->
			let rec f bi =
				let a', (ao, bo) = transition a (ai, bi) in
				match a' with
				| Const (ao', bo') ->
					Some (constant ao'), ao
				| _ ->
					if bo =* bi then Some (loop bo a'), ao
					else f bi
			in f init)

let rec loopq init a =
	make_sf (fun ai ->
			let rec f bi =
				let a', (ao, bo) = transition a (ai, bi) in
				match a' with
				| Const (ao', bo') ->
					Some (constant ao'), ao
				| _ ->
					if bo == bi then Some (loop bo a'), ao
					else f bi
			in f init)

let fix g f =
	let sf' = ref (make_sf (fun _ -> raise Undefined_signal)) in
	let sf = make_sf (fun i -> let sf, o = transition !sf' i in Some sf, o) in
	let r = f sf in
	sf' := g r;
	r

let identity = returnA


(* Event signals *)

let never = Const None

let now v =
	make_sf (fun _ -> Some never, Some v)

let edge v =
	let rec sf = function
		| true -> make_sf (fun i -> Some (sf i), None)
		| false -> make_sf (fun i -> Some (sf i), if i then Some v else None)
	in sf true

let delta f =
	let rec sf init =
		make_sf (fun i ->
			if i != init then Some (sf i), Some (f init i) else None, None)
	in make_sf (fun i -> Some (sf i), None)

let changes =
	let rec sf init =
		make_sf (fun i ->
			if i <>* init then Some (sf i), Some i else None, None)
	in TF (fun i -> Some (sf i), None)

let changesq =
	let rec sf init =
		make_sf (fun i ->
			if i != init then Some (sf i), Some i else None, None)
	in TF (fun i -> Some (sf i), None)

let latch =
	let rec sf init = make_sf (function
		| i, Some _ -> Some (sf i), Some init
		| i, None -> Some (sf i), None)
	in TF (function
		| i, Some _ -> Some (sf i), Some i
		| i, None -> Some (sf i), None)

let once = TF (function
	| Some i -> Some never, Some i
	| None -> None, None)

let rec take = function
	| 0 -> never
	| n ->
		make_sf (function
			| Some i -> Some (take (n - 1)), Some i
			| None -> None, None)

let rec drop = function
	| 0 -> identity
	| n ->
		make_sf (function
			| Some i -> Some (drop (n - 1)), None
			| None -> None, None)


(* Parallel signals *)

let rec par l =
	make_sf (fun i ->
		let l', o = List.split (List.map (fun sf -> transition sf i) l) in
		if List.for_all is_constant l' then
			let o' = List.rev (List.fold_left
				(fun o' -> function
				 | Const ao' -> ao' :: o'
				 | _ -> assert false)
				[] l') in
			Some (constant o'), o
		else Some (par l'), o)


(* Behavior signals *)

let rec hold init =
	make_sf (function Some i -> Some (hold i), i | None -> None, init)

let track = hold

let gate =
	let rec sf init = make_sf (function
		| i, Some _ -> Some (sf i), i
		| _, None -> None, init) in
	TF (function i, _ -> Some (sf i), i)


(* Accumulators *)

let rec accum init =
	make_sf (function
		| Some f -> let o = f init in Some (accum o), Some o
		| None -> None, None)

let rec collect f init =
	make_sf (function
		| Some i -> let o = f init i in Some (collect f o), Some o
		| None -> None, None)


(* Switchers *)


let rec switch sf c =
	make_sf (fun i ->
		match transition c i with
		| c', None ->
			let sf', o = transition sf i in
			if is_constant c' then Some sf', o
			else Some (switch sf' c'), o
		| c', Some sf ->
			let sf', o = transition sf i in
			if is_constant c' then Some sf', o
			else Some (switch sf' c'), o)

(*let rswitch sf =
	switch (arr fst >>> sf) (arr (fun (i, c) ->
		Event.map (fun sf -> arr fst >>> sf) c))*)

let rec rswitch sf =
	make_sf (function
		| i, None ->
			let sf', o = transition sf i in
			Some (rswitch sf'), o
		| i, Some sf ->
			let sf', o = transition sf i in
			Some (rswitch sf'), o)

(*let switch sf c =
	(identity &&& c) >>> rswitch sf*)
	
let rec dswitch (sf: ('a, 'b Event.t) sf) c =
	make_sf (fun i ->
		match transition c i with
		| c', None ->
			let sf', o = transition sf i in
			if is_constant c' then Some sf', o
			else Some (dswitch sf' c'), o
		| c', Some dsf ->
			let _, o = transition sf i
			(* we must transition the new signal function in order to
			   preserve the property of event conservation *)
			and sf', _ = transition dsf i in
			if is_constant c' then Some sf', o
			else Some (dswitch sf' c'), o)

let rec drswitch (sf: ('a, 'b Event.t) sf) =
	make_sf (function
		| i, None ->
			let sf', o = transition sf i in
			Some (drswitch sf'), o
		| i, Some dsf ->
			let _, o = transition sf i
			(* we must transition the new signal function in order to
			   preserve the property of event conservation *)
			and sf', _ = transition dsf i in
			Some (drswitch sf'), o)

let rec srswitch g f init =
	let rec srswitch_r sf o =
		make_sf (function
			| i, None ->
				let sf', o' = transition sf i in
				Some (srswitch_r sf' o'), o'
			| i, Some f ->
				let sf', o' = transition (f (g o)) i in
				Some (srswitch_r sf' o'), o')
	in make_sf (function
		| i, None ->
			let sf', o' = transition (f init) i in
			Some (srswitch_r sf' o'), o'
		| i, Some f ->
			let sf', o' = transition (f init) i in
			Some (srswitch_r sf' o'), o')

let rec rcollapse =
	let rec rcollapse_r cur sf =
		make_sf (function
			| i, cur' when cur' == cur ->
				let sf', o' = transition sf i in
				Some (rcollapse_r cur' sf'), o'
			| i, cur' ->
				let sf', o' = transition cur' i in
				Some (rcollapse_r cur' sf'), o')
	in
	TF (fun (i, cur') ->
		let sf', o' = transition cur' i in
		Some (rcollapse_r cur' sf'), o')

let collapse c =
	(identity &&& c) >>> rcollapse

let rec srcollapse g init =
	let rec srcollapse_r f sf o =
		make_sf (function
			| i, f' when f' == f ->
				let sf', o' = transition sf i in
				Some (srcollapse_r f' sf' o'), o'
			| i, f' ->
				let sf', o' = transition (f' (g o)) i in
				Some (srcollapse_r f' sf' o'), o')
	in make_sf (fun (i, f') ->
		let sf', o' = transition (f' init) i in
		Some (srcollapse_r f' sf' o'), o')

let rec acollapse map =
	make_sf (fun (i, s) ->
		let v', o = transition (List.assoc s map) i in
		Some (acollapse
			(List.map
				(fun (k, v) -> k, if k = s then v' else fst (transition v i))
				map)),
		o)

let const_arr f =
	make_sf (fun i ->
		let stop, o = f i in
		if stop then Some (constant o), o
		else None, o)

let reactimate sf i = transition sf i