import os, glob, re, getopt, sys
import utils
import networkx as nx
import numpy as np
import settings


def saveSnapshot(tempG, threshold, outputFile):
        if threshold>1:
            tempNodeLabels = tempG.nodes()
            adjMatr = nx.to_numpy_matrix(tempG)
            newAdjMatr = np.maximum(adjMatr  - (threshold - 1) * np.ones(adjMatr.shape), np.zeros(adjMatr.shape))
            tempG = nx.relabel_nodes(nx.from_numpy_matrix(newAdjMatr), dict(zip(range(0,tempG.order()), tempNodeLabels)))

        nx.write_edgelist(nx.Graph(tempG), outputFile, data=False)

def splitIntoSnapshots(inputTemporalNetworkFile, deltaT, outputDirSnapshots, threshold = 1, delimiter = ' ', mintT0 = 0, accumulating = False):

    deltaTUnitType = None

    if 'month' in deltaT:
        deltaT = int(re.findall(r'\d+', deltaT)[0]) * 3600 * 24 * 30
        #deltaT = int(re.findall(r'\d+', deltaT)[0])
        #deltaTUnitType = 'month'
    elif 'week' in deltaT:
        deltaT = int(re.findall(r'\d+', deltaT)[0]) * 3600 * 24 * 7
    elif 'day' in deltaT:
        deltaT = int(re.findall(r'\d+', deltaT)[0]) * 3600 * 24
    elif 'hour' in deltaT:
        deltaT = int(re.findall(r'\d+', deltaT)[0]) * 3600
    else:
        deltaT = int(deltaT)

    #if output directory does not exists, create it
    if not os.path.exists(outputDirSnapshots):
        os.makedirs(outputDirSnapshots)

    tempG = nx.MultiDiGraph()

    t0 = None
    prevT = None
    deltaT0Unit = None
    deltaTUnit = None

    i_counter = 1
    with open(inputTemporalNetworkFile, "r") as f:
      for line in f:
        t,u,v = line.split(delimiter)
        t = int(t)
        u = int(u)
        v = int(v)

        if t < mintT0:
            continue

        #remove self-loops
        if u==v:
            continue

        if t0 is None:
            t0 = t
            if deltaTUnitType=='month':
                deltaT0Unit = int(utils.timestamp_to_time(t0, formatString=("%m")))


        tempG.add_edge(u,v)

        if deltaTUnitType=='month':
            deltaTUnit = int(utils.timestamp_to_time(t, formatString=("%m")))

        snapshotName = None
        if deltaTUnitType == 'month' and (deltaTUnit - deltaT0Unit >= deltaT or (deltaTUnit < deltaT0Unit and deltaTUnit + 12 - deltaT0Unit >= deltaT)):
            print '%s. Saving snapshot from %s to %s%s' % (i_counter, utils.timestamp_to_time(t0),utils.timestamp_to_time(t), ' (accumulating edges)' if accumulating else '')
            snapshotName = 'snapshot_%s.txt' % utils.timestamp_to_time(prevT, '%Y-%m')
            deltaT0Unit = int(utils.timestamp_to_time(t0, formatString=("%m")))
        elif deltaTUnitType == None and t - t0 >= deltaT:
            print '%s. Saving snapshot from %s to %s%s' % (i_counter, t0 , t, ' (accumulating edges)' if accumulating else '')
            snapshotName = 'snapshot_%s.txt' % i_counter
            t0 = t

        if snapshotName is not None:
            saveSnapshot(tempG, threshold, outputDirSnapshots + snapshotName)
            i_counter += 1
            if not accumulating:
                tempG.clear()



        prevT = t

    #saving the remaining snapshot
    if (deltaTUnitType == None and t - t0 > 0) or (deltaTUnitType == 'month' and deltaTUnit - deltaT0Unit > 0):#if tempG.order()>0:
        if deltaTUnitType == 'month':
            print '%s. Saving the remaining snapshot from %s to %s%s' % (i_counter, utils.timestamp_to_time(t0),utils.timestamp_to_time(t), ' (accumulating edges)' if accumulating else '')
        else:
            print '%s. Saving the remaining last snapshot from %s to %s%s' % (i_counter, t0,t, ' (accumulating edges)' if accumulating else '')

        snapshotName = 'snapshot_%s.txt' % i_counter
        if deltaTUnitType == 'month':
            snapshotName = 'snapshot_%s.txt' % utils.timestamp_to_time(prevT, '%Y-%m')

        saveSnapshot(tempG, threshold, outputDirSnapshots + snapshotName)


inputSubDir = ''
delimiter = ' '
directed = False
mintT0 = -1
accumulating = False

opts, args = getopt.getopt(sys.argv[1:], "s:d:m:ath", ["subdir", "help"])
for o, a in opts:
    if o in ('-s', '--subdir'):
        inputSubDir = a  + '/'
    elif o == '-d':
        delimiter = a.strip()
        if delimiter == 'tab':
            delimiter = '\t'
    elif o == '-m':
        mintT0 = int(a)
    elif o =='--directed':
        directed = True
    elif o == '-a':
        accumulating = True
    if o in ('-h','--help'):
        print """Usage:
%s network_name delta_t threshold
-s,--subdir      Input/output subdirectory
-d               Delimiter (use 'tab' for '\\t')
-m               Starting time to consider
-a               Accumulate edges across snapshots with time
-h,--help        Print this"""  % sys.argv[0]
        sys.exit(0)

networkName = args[0]
deltaT = args[1]
threshold = int(args[2])

outputDir = settings.BASE_INPUT_DIR + inputSubDir + networkName + '/snapshots/snapshots_%s_%s%s/' % (deltaT, threshold, '_acc' if accumulating else '')

splitIntoSnapshots(settings.BASE_INPUT_DIR + inputSubDir + networkName + '/' + networkName + '.txt', deltaT, outputDir, threshold, delimiter, mintT0 = mintT0, accumulating = accumulating)