from collections import defaultdict
from itertools import product
import grammarette

one_d = {
	(0,): "buviko",
	(0,): "buviko",
	(0,): "buviko",
	(0,): "buviko",
	(1,): "zetiko",
	(1,): "zetiko",
	(1,): "zetiko",
	(1,): "zetiko",
	(2,): "gafiko",
	(2,): "gafiko",
	(2,): "gafiko",
	(2,): "gafiko",
	(3,): "wopiko",
	(3,): "wopiko",
	(3,): "wopiko",
	(3,): "wopiko",
}
transparent = {
	(0, 0): "buviko",
	(0, 1): "buviko",
	(0, 2): "buviko",
	(0, 3): "buviko",
	(1, 0): "zetiko",
	(1, 1): "zetiko",
	(1, 2): "zetiko",
	(1, 3): "zetiko",
	(2, 0): "gafiko",
	(2, 1): "gafiko",
	(2, 2): "gafiko",
	(2, 3): "gafiko",
	(3, 0): "wopiko",
	(3, 1): "wopiko",
	(3, 2): "wopiko",
	(3, 3): "wopiko",
}
holistic = {
	(0, 0): "buvichoe",
	(0, 1): "buvikoh",
	(0, 2): "buvicow",
	(0, 3): "buvycow",
	(1, 0): "zeteekow",
	(1, 1): "zeteecoh",
	(1, 2): "zeticoe",
	(1, 3): "zeticoh",
	(2, 0): "gafeechow",
	(2, 1): "gafeekoe",
	(2, 2): "gafykoe",
	(2, 3): "gafychoe",
	(3, 0): "wopykow",
	(3, 1): "wopeecoe",
	(3, 2): "wopykoh",
	(3, 3): "wopikow",
}
redundant = {
	(0, 0): "buvikoh",
	(0, 1): "buvikoh",
	(0, 2): "buvikoh",
	(0, 3): "buvikoh",
	(1, 0): "zeteeco",
	(1, 1): "zeteeco",
	(1, 2): "zeteeco",
	(1, 3): "zeteeco",
	(2, 0): "gafykoe",
	(2, 1): "gafykoe",
	(2, 2): "gafykoe",
	(2, 3): "gafykoe",
	(3, 0): "wopycow",
	(3, 1): "wopycow",
	(3, 2): "wopycow",
	(3, 3): "wopycow",
}
expressive = {
	(0, 0): "buvikoh",
	(0, 1): "buveeco",
	(0, 2): "buvykoe",
	(0, 3): "buvycow",
	(1, 0): "zetikoh",
	(1, 1): "zeteeco",
	(1, 2): "zetykoe",
	(1, 3): "zetycow",
	(2, 0): "gafikoh",
	(2, 1): "gafeeco",
	(2, 2): "gafykoe",
	(2, 3): "gafycow",
	(3, 0): "wopikoh",
	(3, 1): "wopeeco",
	(3, 2): "wopykoe",
	(3, 3): "wopycow",
}
expressive_letter = {
	(0, 0): "buvikop",
	(0, 1): "buvikog",
	(0, 2): "buvikob",
	(0, 3): "buvikoy",
	(1, 0): "zetikop",
	(1, 1): "zetikog",
	(1, 2): "zetikob",
	(1, 3): "zetikoy",
	(2, 0): "gafikop",
	(2, 1): "gafikog",
	(2, 2): "gafikob",
	(2, 3): "gafikoy",
	(3, 0): "wopikop",
	(3, 1): "wopikog",
	(3, 2): "wopikob",
	(3, 3): "wopikoy",
}
expressive_length = {
	(0, 0): "buviko",
	(0, 1): "buvikoo",
	(0, 2): "buvikooo",
	(0, 3): "buvikoooo",
	(1, 0): "zetiko",
	(1, 1): "zetikoo",
	(1, 2): "zetikooo",
	(1, 3): "zetikoooo",
	(2, 0): "gafiko",
	(2, 1): "gafikoo",
	(2, 2): "gafikooo",
	(2, 3): "gafikoooo",
	(3, 0): "wopiko",
	(3, 1): "wopikoo",
	(3, 2): "wopikooo",
	(3, 3): "wopikoooo",
}
redundant_expressive = {
	(0, 0): "buvikoe",
	(0, 1): "buvikoh",
	(0, 2): "buvikoe",
	(0, 3): "buvichoe",
	(1, 0): "zeteekoe",
	(1, 1): "zeteekoh",
	(1, 2): "zeteekoe",
	(1, 3): "zeteechoe",
	(2, 0): "gafykoe",
	(2, 1): "gafykoh",
	(2, 2): "gafykoe",
	(2, 3): "gaffychoe",
	(3, 0): "wopykoe",
	(3, 1): "wopykoh",
	(3, 2): "wopykoe",
	(3, 3): "wopychoe",
}
kcs_example1 = { # (color, movement, shape)
	(0, 0, 0): "tuge",
	(0, 0, 1): "tuge",
	(0, 0, 2): "tuge",
	(0, 1, 0): "tupim",
	(0, 1, 1): "miniku",
	(0, 1, 2): "tupin",
	(0, 2, 0): "poi",
	(0, 2, 1): "poi",
	(0, 2, 2): "poi",
	(1, 0, 0): "tuge",
	(1, 0, 1): "tuge",
	(1, 0, 2): "tuge",
	(1, 1, 0): "tupim",
	(1, 1, 1): "miniku",
	(1, 1, 2): "tupin",
	(1, 2, 0): "poi",
	(1, 2, 1): "poi",
	(1, 2, 2): "poi",
	(2, 0, 0): "tuge",
	(2, 0, 1): "tuge",
	(2, 0, 2): "tuge",
	(2, 1, 0): "tupim",
	(2, 1, 1): "miniku",
	(2, 1, 2): "tupin",
	(2, 2, 0): "poi",
	(2, 2, 1): "poi",
	(2, 2, 2): "poi",
}
kcs_example2 = { # (color, movement, shape)
	(0, 0, 0): "nereki",
	(0, 0, 1): "neheki",
	(0, 0, 2): "nekeki",
	(0, 1, 0): "nereplo",
	(0, 1, 1): "nehoplo",
	(0, 1, 2): "nekiplo",
	(0, 2, 0): "nepilu",
	(0, 2, 1): "nehopilu",
	(0, 2, 2): "nekipilu",
	(1, 0, 0): "lereki",
	(1, 0, 1): "lahoki",
	(1, 0, 2): "lakeki",
	(1, 1, 0): "laneplo",
	(1, 1, 1): "lahoplo",
	(1, 1, 2): "lakiplo",
	(1, 2, 0): "lanepilu",
	(1, 2, 1): "lahopilu",
	(1, 2, 2): "lakipilu",
	(2, 0, 0): "renana",
	(2, 0, 1): "reneki",
	(2, 0, 2): "raheki",
	(2, 1, 0): "replo",
	(2, 1, 1): "rehoplo",
	(2, 1, 2): "rahoplo",
	(2, 2, 0): "repilu",
	(2, 2, 1): "rehopilu",
	(2, 2, 2): "rahopilu",
}
ktcs_example1 = {
	(0, 0): "pihino",
	(0, 1): "nemone",
	(0, 2): "piga",
	(0, 3): "kawake",
	(1, 0): "kapa",
	(1, 1): "gakho",
	(1, 2): "wuwele",
	(1, 3): "nepi",
	(2, 0): "newhomo",
	(2, 1): "kamone",
	(2, 2): "gaku",
	(2, 3): "hokako",
}
ktcs_example2 = {
	(0, 0): "egewawu",
	(0, 1): "egewawa",
	(0, 2): "egewuwu",
	(0, 3): "ege",
	(1, 0): "mega",
	(1, 1): "megawawa",
	(1, 2): "megawuwu",
	(1, 3): "wulagi",
	(2, 0): "gamanewawu",
	(2, 1): "gamanewawa",
	(2, 2): "gamanewuwu",
	(2, 3): "gamane",
}

test_set = [
	(one_d, (4,), '0buv|1zet|2gaf|3wop+?iko'),
	(transparent, (4, 4), '0?buv|1?zet|2?gaf|3?wop+??iko'),
	(holistic, (4, 4), '0?buv|1?zet|2?gaf|3?wop+00ichoe|01ikoh|02icow|03ycow|10eekow|11eecoh|12icoe|13icoh|20eechow|21eekoe|22ykoe|23ychoe|30ykow|31eecoe|32ykoh|33ikow'),
	(redundant, (4, 4), '0?buvikoh|1?zeteeco|2?gafykoe|3?wopycow'),
	(expressive, (4, 4), '0?buv|1?zet|2?gaf|3?wop+?0ikoh|?1eeco|?2ykoe|?3ycow'),
	(expressive_letter, (4, 4), '0?buv|1?zet|2?gaf|3?wop+??iko+?0p|?1g|?2b|?3y'),
	(expressive_length, (4, 4), '0?buv|1?zet|2?gaf|3?wop+??iko+?1o|?2oo|?3ooo'),
	(redundant_expressive, (4, 4), '0?buvi|1?zetee|23gaffy|2?gafy|3?wopy+?3ch|??k+?1oh|??oe'),
	(kcs_example1, (3, 3, 3), '?0?tuge|?10tupim|?11miniku|?12tupin|?2?poi'),
	(kcs_example2, (3, 3, 3), '001nehe|002neke|020ne|0?0nere|0?1neho|0?2neki|100lere|102lake|1?0lane|1?1laho|1?2laki|200ren|201rene|202rahe|2?0re|2?1reho|2?2raho+200ana|?0?ki|?1?plo|?2?pilu'),
	(ktcs_example1, (3, 4), '00pihino|01nemone|02piga|03kawake|10kapa|11gakho|12wuwele|13nepi|20newhomo|21kamone|22gaku|23hokako'),
	(ktcs_example2, (3, 4), '0?ege|10meg|13wulagi|1?mega|2?gamane+10a|?0wawu|?1wawa|?2wuwu'),
]

def test_grammar(grammar, lexicon):
	signal_to_meanings = defaultdict(set)
	for meaning, expected_signal in lexicon.items():
		signal_to_meanings[expected_signal].add(meaning)
		produced_signal = grammar.produce(meaning)
		if expected_signal != produced_signal:
			print(f'❗️ Production error for {meaning}: Expected "{expected_signal}", produced "{produced_signal}"')
	for signal, expected_meanings in signal_to_meanings.items():
		inferred_meanings = grammar.comprehend(signal)
		if inferred_meanings != expected_meanings:
			print(f'❗️ Comprehension error for "{signal}": Expected {expected_meanings}, inferred {inferred_meanings}')

for test_lexicon, dims, expected_grammar in test_set:

	G = grammarette.induce(test_lexicon, dims)
	print(G)
	print(G.codelength)

	assert G.grammar == expected_grammar
	test_grammar(G, test_lexicon)
