import json
import math
import pdb
import random
import numpy as np
import sys
import time
sys.path.extend(['.', '..'])

from itertools import chain, repeat, izip
from collections import defaultdict
from operator import mul, and_, or_
from scipy.optimize import fsolve
from matplotlib import pyplot as plt

from ..util import *
from ..bottomup.bounding_box import *
from ..bottomup.cluster import *

from frontier import *
from rangemerger import RangeMerger2

_logger = get_logger()


    
 
class StreamRangeMerger(RangeMerger2):
  def __init__(self, *args, **kwargs):
    super(StreamRangeMerger, self).__init__(*args, **kwargs)

    self.valid_cluster_f = kwargs.get('valid_cluster_f', lambda c: True)

    # idx -> clusters to expand  -- different than clusters on frontier!!
    self.tasks = defaultdict(list)

    # all values for each dimension
    self.all_cont_vals = defaultdict(set) # idx -> values

    # attribute name -> { attr val -> [sum of influence value at c=0.1, count] }
    self.all_disc_vals = defaultdict(lambda: defaultdict(lambda: [0,0])) 

    # name -> { val -> # times failed }
    self.failed_disc_vals = defaultdict(lambda: defaultdict(lambda:0))

    # stores the frontier after each iteration
    self.added = set()
    self.seen = set()
    self.frontiers = []
    self.adj_graph = None


    self.K = 2
    self.nblocks = 50

    if len(self.learner.full_table) < 40000:
      self.K = 2
      self.nblocks = 60
    if len(self.learner.full_table) < 10000:
      self.nblocks = 100

    self.get_frontier = CheapFrontier(self.c_range, K=self.K, nblocks=self.nblocks, learner=self.learner)
    self.get_frontier.stats = self.stats


    if self.DEBUG:
      self.renderer = InfRenderer('/tmp/merger.pdf', c_range=self.c_range)

  def close(self):
    if self.DEBUG:
      self.renderer.close()

  def get_frontier_obj(self, version):
    while version >= len(self.frontiers):
      frontier = CheapFrontier(self.c_range, K=self.K, nblocks=self.blocks, learner=self.learner)
      frontier.stats = self.stats
      self.frontiers.append(frontier)
    return self.frontiers[version]

  @property
  def frontier_iter(self):
    return list(self.frontiers)


  @instrument
  def setup_stats(self, clusters):
    all_inf = lambda l: all([abs(v) == float('inf') for v in l])
    clusters = filter(lambda c: c.bound_hash not in self.added, clusters)
    clusters = filter(lambda c: not all_inf(c.inf_state[0]), clusters)
    clusters = filter(lambda c: len(c.inf_state[2]) == 0 or not all_inf(c.inf_state[2]), clusters)
    self.added.update([c.bound_hash for c in clusters])

    super(StreamRangeMerger, self).setup_stats(clusters)

    start = time.time()
    if not self.adj_graph:
      self.adj_graph = self.make_adjacency([], True)
    self.adj_graph.insert(clusters)
    self.adj_graph.sync()
    self.stats['adj_sync'][0] += time.time() - start
    self.stats['adj_sync'][1] += 1

    for c in clusters:
      for idx in xrange(len(c.cols)):
        self.all_cont_vals[idx].add(c.bbox[0][idx])
        self.all_cont_vals[idx].add(c.bbox[1][idx])
      for disc, vals in c.discretes.iteritems():
        if len(vals) < 3:
          vals = [(v,) for v in vals]
        else:
          vals = [tuple(vals)]
        for v in vals:
          self.all_disc_vals[disc][v][0] += c.inf_func(0.1)
          self.all_disc_vals[disc][v][1] += 1
        #self.all_disc_vals[disc].update(vals)

    return clusters

  @instrument
  def best_so_far(self, prune=False):
    clusters = set()
    for frontier in self.frontier_iter:
      clusters.update(frontier.frontier)

    if prune:
      for c in clusters:
        c.c_range = list(self.c_range)
      clusters = self.get_frontier(clusters)[0]
      clusters = filter(lambda c: r_vol(c.c_range), clusters)

    if self.DEBUG:
      self.renderer.new_page()
      self.renderer.set_title('best so far')
      self.renderer.plot_active_inf_curves(clusters)
    return clusters

  @instrument
  def best_at_c(self, c_val, K=6):
    clusters = set()
    for frontier in self.frontier_iter:
      clusters.update(frontier.seen_clusters)
    rm_dups(clusters, key=lambda c: str(c.rule))

    clusters = sorted(clusters, key=lambda c: c.inf_func(c_val), reverse=True)[:K]
    return clusters

  @instrument
  def add_clusters(self, clusters, idx=0):
    """
    Return list of new clusters that are on the frontier
    """
    if not clusters: return []

    if self.DEBUG:
      print "add_clusters"
      self.print_clusters(clusters)
      self.renderer.new_page()
      self.renderer.set_title("add_clusters %d clusters" % len(clusters))
      for f in self.frontier_iter:
        self.renderer.plot_inf_curves(f.frontier, color='grey') 
      self.renderer.plot_inf_curves(clusters, color='green')

    clusters = self.setup_stats(clusters)
    base_frontier = self.get_frontier_obj(idx)
    clusters, _ = base_frontier.update(clusters)

    if self.DEBUG:
      print "base_frontier"
      self.print_clusters(clusters)
      self.renderer.plot_active_inf_curves(clusters, color='red')

    # clear out current tasks
    self.tasks[idx] = filter(base_frontier.__contains__, self.tasks[idx])
    self.tasks[idx].extend(clusters)

    # remove non-frontier-based expansions from future expansion
    for tidx in self.tasks.keys():
      if tidx <= idx: continue
      checker = lambda c: not any(map(base_frontier.__contains__, c.ancestors))
      self.tasks[tidx] = filter(checker, self.tasks[tidx])

    if clusters:
      _logger.debug("merger:\tadded %d clusters\t%d tasks left", len(clusters), self.ntasks)
    return clusters


  @property
  def ntasks(self):
    if len(self.tasks) == 0: return 0
    return sum(map(len, self.tasks.values()))

  def has_next_task(self):
    if not self.tasks: return False
    return self.ntasks > 0

  def next_tasks(self, n=1):
    ret = []
    for tkey in reversed(self.tasks.keys()):
      tasks = self.tasks[tkey]
      while len(ret) < n and tasks:
        idx = random.randint(0, len(tasks)-1)
        ret.append((idx, tasks.pop(idx)))
    return ret



  @instrument
  def __call__(self, n=2):
    """
    Return any successfully expanded clusters (improvements)
    """
    nmerged = self.nmerged
    start = time.time()
    tasks = self.next_tasks(n)
    improvements = set()

    for idx, cluster in tasks:
      cur_frontier = self.get_frontier_obj(idx)
      next_frontier = self.get_frontier_obj(idx+1)
      new_clusters = self.run_task(idx, cluster, cur_frontier, next_frontier)

      debug = self.DEBUG
      self.DEBUG = False
      self.add_clusters(new_clusters, idx+1)
      self.DEBUG = debug
      improvements.update(new_clusters)

    _logger.debug("merger\ttook %.1f sec\t%d improved\t%d tried\t%d tasks left", 
        time.time()-start, len(improvements), (self.nmerged-nmerged), self.ntasks)
    return improvements


  def run_task(self, idx, cluster, cur_frontier, next_frontier):
    if not (idx == 0 or self.valid_cluster_f(cluster)):
      _logger.debug("merger\tbelow thresh skipping\t %s" % cluster)
      return []

    if self.DEBUG:
      self.renderer.new_page()
      self.renderer.set_title("expand %s" % str(cluster.rule))
      self.renderer.plot_inf_curves([cluster], color='grey') 

    self.rejected_disc_vals = defaultdict(list)
    self.rejected_cont_vals = defaultdict(set)
    expanded = self.greedy_expansion(cluster, self.seen, idx, cur_frontier)
    expanded = [c for c in expanded if c.bound_hash != cluster.bound_hash]

    if self.DEBUG:
      self.renderer.plot_inf_curves(expanded, color='green')

    cur_expanded, rms = cur_frontier.update(expanded)
    next_expanded, rms2 = next_frontier.update(cur_expanded)
    f = lambda c: c.bound_hash != cluster.bound_hash
    improved_clusters = set(filter(f, next_expanded))

    to_hash = lambda cs: set([c.bound_hash for c in cs])
    exp_bounds = to_hash(expanded)
    cur_bounds = to_hash(cur_expanded)
    next_bounds = to_hash(next_expanded)

    for c in chain(cur_expanded, rms):
      _logger.debug("merger\texpanded\tcur_idx(%s)\tnext_idx(%s)\t%.3f-%.3f\t%s", 
          (c.bound_hash in exp_bounds), 
          (c.bound_hash in next_bounds), 
          c.c_range[0], c.c_range[1],
          c.rule.simplify())
    if self.DEBUG:
      self.renderer.plot_active_inf_curves(cur_frontier.frontier, color='blue')
      self.renderer.plot_active_inf_curves(next_frontier.frontier, color='red')

    return improved_clusters


  @instrument
  def dims_to_expand(self, cluster, seen, version=None):
    for idx in xrange(len(cluster.cols)):
      vals = np.array(list(self.all_cont_vals[idx]))
      smaller = vals[(vals < cluster.bbox[0][idx])]
      bigger =  vals[(vals > cluster.bbox[1][idx])]
      yield idx, 'dec', smaller.tolist()
      yield idx, 'inc', bigger.tolist()

    for name, vals in cluster.discretes.iteritems():
      ret = []
      maxval = (len(vals) > 1) and max(vals) or None
      vals2infs = self.all_disc_vals[name].items()
      vals2infs.sort(key=lambda p: p[1][0] / float(p[1][1]+1.), reverse=True)

      for disc_vals, score in vals2infs:
        subset = set(disc_vals).difference(vals)
        subset.difference_update([v for v in subset if self.failed_disc_vals[name][str(v)] > 1])
        if maxval:
          subset = set(filter(lambda v: v >= maxval, subset))
        ret.append(subset)
      ret = filter(bool, ret)

      if ret:
        yield name, 'disc', ret
        return
        p = np.arange(len(ret), 0, -1).astype(float)
        p /= p.sum()
        ret = np.random.choice(ret, min(len(ret), 10), p=p, replace=False)
        yield name, 'disc', ret

  @instrument
  def check_direction(self, cluster, dim, direction, vals):
    key = cluster.bound_hash
    if direction == 'disc':
      for subset in self.rejected_disc_vals[dim]:
        if subset.issubset(vals): 
          return []
    if direction == 'inc':
      cont_vals = self.rejected_cont_vals[(dim, direction)]
      if cont_vals:
        vals = filter(lambda v: v > max(cont_vals), vals)
    if direction == 'dec':
      cont_vals = self.rejected_cont_vals[(dim, direction)]
      if cont_vals:
        vals = filter(lambda v: v < min(cont_vals), vals)
    return vals

  @instrument
  def update_rejected_directions(self, cluster, dim, direction, val):
    if direction == 'disc':
      if not hasattr(val, '__iter__'):
        val = [val]
      for v in list(val):
        self.rejected_disc_vals[dim].append(set([v]))
        self.failed_disc_vals[dim][str(v)] += 1
    if direction == 'inc':
      self.rejected_cont_vals[(dim, direction)].add(round(val, 1))
    if direction == 'dec':
      self.rejected_cont_vals[(dim, direction)].add(round(val, 1))

  @instrument
  def greedy_expansion(self, cluster, seen, version=None, frontier=None):
    _logger.debug("merger\tgreedy_expand\t%s", cluster.rule.simplify())
    if frontier is None:
      frontier = CheapFrontier(self.c_range, K=1, nblocks=15, learner=self.learner)
      frontier.stats = self.stats
      frontier.update([cluster])

    cols = cluster.cols
    for dim, direction, vals in self.dims_to_expand(cluster, seen, version=version):
      if len(vals) == 0: continue
      attrname = isinstance(dim, basestring) and dim or cols[dim]
      vals = self.check_direction(cluster, dim, direction, vals)
      realvals = self.pick_expansion_vals(cluster, dim, direction, vals)

      nfails = 0
      for v in realvals:
        tmp = None
        if direction == 'inc':
          tmp = self.dim_merge(cluster, dim, None, v, seen)
        elif direction == 'dec':
          tmp = self.dim_merge(cluster, dim, v, None, seen)
        else:
          tmp = self.disc_merge(cluster, dim, v)

        if not tmp: 
          _logger.debug("merger\tnoexpand\t%s\t%s\t%s options", attrname[:15], direction, len(vals))
          continue

        improvements = frontier.improvement(tmp)
        if improvements.max() > 0:
          print str(tmp)
          print "\t", [round(v,2) for v in improvements]
        frontier.update([tmp])
        isbetter = tmp in frontier
        _logger.debug("merger\tcand\t%s\t%s\t%s\t%s", attrname[:15], direction, isbetter, v)
 
        seen.add(tmp.bound_hash)
        if not isbetter:
          self.update_rejected_directions(cluster, dim, direction, v)
          if direction != 'disc':
            nfails += 1
            if nfails > 1:
              break
        if direction != 'disc':
          cluster = tmp

    return frontier.frontier


class PartitionedStreamRangeMerger(StreamRangeMerger):
  """
  Partitions the merger based on user defined labels so that frontier curves from 
  one partition do not suppress curves in another partition

  MR labels based on dimensionality
  BDT labels as leaf/non-leaf
  """

  def __init__(self, *args, **kwargs):
    super(PartitionedStreamRangeMerger, self).__init__(*args, **kwargs)
    self.frontiers = defaultdict(list)
    self.tasks = defaultdict(list)

  def get_frontier_obj(self, version, partitionkey):
    frontiers = self.frontiers[partitionkey]
    while version >= len(frontiers):
      frontier = CheapFrontier(self.c_range, K=self.K, nblocks=self.nblocks, learner=self.learner)
      frontier.stats = self.stats
      frontiers.append(frontier)
    return frontiers[version]
  
  @property
  def frontier_iter(self):
    return chain(*self.frontiers.values())

  @instrument
  def add_clusters(self, 
      clusters, idx=0, partitionkey=None, skip_frontier=False):
    """
    Return list of new clusters that are on the frontier
    """
    if partitionkey is None:
      raise RuntimeError('addclusters partitionkey cannot be none')

    if not clusters: return []

    print "add %d clusters" % len(clusters)
    if self.DEBUG:
      self.renderer.new_page()
      self.renderer.set_title("add_clusters %d clusters" % len(clusters))
      for f in self.frontier_iter:
        self.renderer.plot_inf_curves(f.frontier, color='grey') 
      self.renderer.plot_inf_curves(clusters, color='green')

    nclusters = len(clusters)
    clusters = self.setup_stats(clusters)
    frontier = self.get_frontier_obj(idx, partitionkey)
    if not skip_frontier:
      clusters, _ = frontier.update(clusters)

    # XXX: new cluster should be better than _all_ frontiers
    #for f in self.frontier_iter:
    #clusters, _ = f.update(clusters)

    if not clusters: 
      return clusters

    if self.DEBUG and not skip_frontier:
      print "base_frontier"
      self.print_clusters(clusters)

    if self.DEBUG:
      self.renderer.plot_active_inf_curves(clusters, color='red')

    # clear out current tasks
    tkey = (partitionkey, idx)
    self.tasks[tkey] = filter(frontier.__contains__, self.tasks[tkey])
    self.tasks[tkey].extend(clusters)

    # remove non-frontier-based expansions from future expansion
    for (pkey, tidx) in self.tasks.keys():
      if pkey != partitionkey: continue
      if tidx <= idx: continue
      checker = lambda c: not any(map(frontier.__contains__, c.ancestors))
      self.tasks[tkey] = filter(checker, self.tasks[tkey])

    _logger.debug("merger\t%s\tadd %d of %d clusters\t%d idx\t%d tasks left", partitionkey, len(clusters), nclusters, idx, self.ntasks)
    return clusters


  def next_tasks(self, n=1):
    ret = []
    for tkey in reversed(self.tasks.keys()):
      if len(ret) >= n: break
      tasks = self.tasks[tkey]
      ntasks = len(tasks)
      if not ntasks: continue
      idxs = np.random.choice(ntasks, min(ntasks, n-len(ret)), replace=False).tolist()
      for idx in sorted(idxs, reverse=True):
        ret.append((tkey[0], tkey[1], tasks.pop(idx)))
    return ret


  def __call__(self, n=2):
    nmerged = self.nmerged
    start = time.time()
    tasks = self.next_tasks(n)
    improvements = set()

    for pkey, idx, cluster in tasks:
      cur_frontier = self.get_frontier_obj(idx, pkey)
      next_frontier = self.get_frontier_obj(idx+1, pkey)
      new_clusters = self.run_task(idx, cluster, cur_frontier, next_frontier)

      self.add_clusters(new_clusters, idx=idx+1, partitionkey=pkey, skip_frontier=True)
      improvements.update(new_clusters)

      _logger.debug("merger\t%s\ttook %.1f sec\t%d improved\t%d tried\t%d tasks left", 
          pkey, time.time()-start, len(improvements), (self.nmerged-nmerged), self.ntasks)
    return improvements




