from loader import Loader
from collections import defaultdict
from typing import Dict, List, Tuple, Set
import numpy as np
import graphviz

DIR = "/home/velocitatem/Documents/Projects/PHANTOM/experiments/collected_data/"

class BehaviorModel:
    def __init__(self, src_dir: str = DIR):
        self.loader = Loader(src_dir)
        self.data = self.loader.get_data()
        self.entries, self.num_entries = self.loader.get_entries()
        self.mdp = None

    def _state_repr(self, evt) -> str:
        p = evt.value.payload
        return f"{p.page or 'unk'}|{p.productId or 'none'}|{p.eventName}"

    def _extract_sessions(self):
        # transform raw events into sequential state trajectories per session
        trajectories = []
        for sid, evts in self.data.items():
            if len(evts) < 2: continue
            states = [self._state_repr(e) for e in sorted(evts, key=lambda x: x.timestamp)]
            trajectories.append(states)
        return trajectories

    def _calc_transitions(self, trajectories: List[List[str]]) -> Tuple[Dict, Set]:
        trans = defaultdict(lambda: defaultdict(int))
        states = set()
        for traj in trajectories:
            for i in range(len(traj) - 1):
                s, s_next = traj[i], traj[i+1]
                trans[s][s_next] += 1
                states.update([s, s_next])
        return trans, states

    def _calc_rewards(self, trajectories: List[List[str]]) -> Dict:
        # reward based on session progression depth
        rwd = defaultdict(list)
        for traj in trajectories:
            n = len(traj)
            for i, s in enumerate(traj):
                rwd[s].append(i / n)
        return rwd

    def _normalize_trans(self, counts: Dict) -> Dict:
        return {s: {s_n: cnt/sum(nxt.values()) for s_n, cnt in nxt.items()}
                for s, nxt in counts.items()}

    def build_MDP(self) -> Dict:
        trajs = self._extract_sessions()
        trans_cnt, states = self._calc_transitions(trajs)
        trans_prob = self._normalize_trans(trans_cnt)
        state_rwd = self._calc_rewards(trajs)
        state_val = {s: np.mean(r) for s, r in state_rwd.items()}

        self.mdp = {
            'states': sorted(list(states)),
            'num_states': len(states),
            'transitions': trans_prob,
            'state_values': state_val,
            'state_rewards': state_rwd,
            'trans_counts': trans_cnt,
        }
        return self.mdp

    def transition_prob(self, s: str, s_next: str) -> float:
        if not self.mdp: raise ValueError("build MDP first")
        return self.mdp['transitions'].get(s, {}).get(s_next, 0.0)

    def state_value(self, s: str) -> float:
        if not self.mdp: raise ValueError("build MDP first")
        return self.mdp['state_values'].get(s, 0.0)

    def sample_traj(self, start: str, max_len: int = 50) -> List[str]:
        if not self.mdp: raise ValueError("build MDP first")
        path = [start]
        curr = start
        for _ in range(max_len):
            nxt = self.mdp['transitions'].get(curr, {})
            if not nxt: break
            curr = np.random.choice(list(nxt.keys()), p=list(nxt.values()))
            path.append(curr)
        return path

def visualize_mdp(model: BehaviorModel, threshold: float = 0.05, output: str = "mdp_graph", fmt: str = "svg", view: bool = False, export_dot: bool = False):
    """visualize MDP as directed graph using graphviz, aggregated by event type"""
    if not model.mdp: raise ValueError("build MDP first")

    # aggregate transitions by event type
    evt_trans = defaultdict(lambda: defaultdict(float))
    for s, trans in model.mdp['transitions'].items():
        evt_src = s.split('|')[2]
        for s_next, prob in trans.items():
            evt_dst = s_next.split('|')[2]
            evt_trans[evt_src][evt_dst] += prob

    # normalize aggregated transitions
    for evt_src in evt_trans:
        total = sum(evt_trans[evt_src].values())
        if total > 0:
            for evt_dst in evt_trans[evt_src]:
                evt_trans[evt_src][evt_dst] /= total

    g = graphviz.Digraph(format=fmt)
    g.attr(rankdir='LR', size='30')
    g.attr('node', shape='circle', width='1', height='1')

    # collect all event types
    events = set(evt_trans.keys())
    for trans in evt_trans.values():
        events.update(trans.keys())

    # add nodes for each event type
    for evt in events:
        g.node(evt)

    # add edges above threshold
    for evt_src in evt_trans:
        for evt_dst, prob in evt_trans[evt_src].items():
            if prob > threshold:
                g.edge(evt_src, evt_dst, label=f'{prob:.2f}')

    g.render(output, view=view, cleanup=True)
    print(f"Saved MDP graph to {output}.{fmt}")

    if export_dot:
        dot_file = f"{output}.dot"
        with open(dot_file, 'w') as f:
            f.write(g.source)
        print(f"Exported DOT source to {dot_file}")

    return g

if __name__ == "__main__":
    model = BehaviorModel(DIR)
    mdp = model.build_MDP()
    print(f"Built MDP: {mdp['num_states']} states, {sum(len(t) for t in mdp['transitions'].values())} transitions")
    if not mdp['states']:
        print("No states found")
        exit(1)
    visualize_mdp(model, threshold=0.05, output="mdp_viz", fmt="pdf", export_dot=True)