;; open Assert
;; open Map.ListMap

type student = string

type major = string

type school = SEAS | WHARTON | SAS | NURSING 

(* A type abbreviation for finite maps from students to majors. *) 
type roster = (student, major) map

(* What school is each major in? *)
let to_school (m : major) : school = begin match m with 
	| "CIS" | "DMD" | "EE" | "BE" | "CMPE" | "SSE" 
	| "MSE" | "CD" | "MEAM" | "MKSE"     -> SEAS
	| "FNCE" | "OPIM" | "STAT" | "MGMT" | "WUNG"  -> WHARTON
	| "LING" | "ECON" | "MATH"  | "PHYS" | "ENGL"
	| "PPE" | "ENVS" | "UNDC"  -> SAS
	| _ -> failwith "Unknwon major"
 end

(* Is a particular major in a given school? *)
let is_engr    (m : major) : bool = to_school m = SEAS
let is_wharton (m : major) : bool = to_school m = WHARTON
let is_sas     (m : major) : bool = to_school m = SAS
let is_nursing (m : major) : bool = to_school m = NURSING

let test_roster = from_list 
  [ ("Ada", "CIS"); ("Brad", "DMD"); 
    ("Tom", "STAT"); ("Noam", "LING") ] 
		

(* ------ Printing -------- *)
let string_of_school (m : school) : string = 
	begin match m with
		| SEAS -> "SEAS" | WHARTON -> "WHARTON" 
		| SAS ->  "SAS"  | NURSING -> "NURSING"
  end

let string_of_roster (r : roster) : string =
	let rec aux (l : (student * major) list) : string = 
  	begin match l with
		| [] -> ""
		| [ (k,v) ] -> k ^ "->" ^  v		
		| (k,v)::t -> k ^ "->" ^ v ^ ", " ^ aux t
	  end
	in aux (Map.ListMap.bindings r)
		
;; print_endline "Test roster"
;; print_endline (string_of_roster test_roster)
(* ------ End printing --------- *)		

let only (f : major -> bool) (m:roster) : roster =
	let rec aux (l : (student * major) list) : (student * major) list = 
		begin match l with 
			| [] -> []
			| (s,m)::t -> if f m then (s,m) :: aux t else aux t
		end
	in
  	from_list (aux (bindings m))

(* Filter the roster so that only engineering students remain. *)					     
let only_engr (m : roster) : roster = 
	only is_engr m
		
(* Filter the roster so that only wharton students remain. *)				
let only_wharton (m : roster) : roster = 
	only is_wharton m
  
	
(* Filter the roster so that only SAS students remain. *)
let only_sas (m : roster) : roster =
	only is_sas m 

(* Filter the roster so that only nursing students remain. *)
let only_nursing (m : roster) : roster = 
	 only is_nursing m 

(* -------------------- test cases ---------------------- *)

let test () : bool =
  (only_engr test_roster) = from_list [ ("Ada", "CIS");("Brad", "DMD") ] 
;; run_test "only_engr test_roster" test 

let test () : bool =  
  (only_wharton test_roster) = from_list [ ("Tom", "STAT") ] 
;; run_test "only_wharton test_roster" test

let test () : bool =
  only_sas test_roster = from_list [ ("Noam", "LING") ] 
;; run_test "only_sas test_roster" test

let test () : bool = 
	only_nursing test_roster = empty
;; run_test "only nursing students"	test

;; print_endline "End of majors.ml"