import os
import os.path
import subprocess
import time
import re
import magic
import zipfile
import rarfile

FILE_NAME_RE = '^([^@]+@[^@]+)\.([0-9]+)(\.[a-zA-Z0-9]+)?$'
FILE_NAME_MATCHER = re.compile(FILE_NAME_RE)
FileType = magic.Magic()
USED_ARCHIVE_DIR = '__inflated_archive_files'

OSX_JUNKS = ['__MACOSX', '.DS_Store']
OSX_JUNK_DIR = '__purged_osx_junks'

def sample_script():
    init()
    moved, notmoved = move_all()
    print 'not moved:', notmoved
    raw_input('Please check the not moved')
    purge_osx_junks('.')
    for f in moved:
        flattern_dir(f)

def init():
    ensure_dir(USED_ARCHIVE_DIR, OSX_JUNK_DIR)

def split_file_name(name):
    """
    split a file name to three parts [email, version, extension]
    """
    matched = FILE_NAME_MATCHER.match(name)
    if matched:
        return matched.group(1, 2, 3)
    else:
        print 'Malformed file name', name
        return None

def ensure_dir(*dirs):
    """
    ensures that all the directories in dirs are created if not already exist
    if you have two directories "a" and "a/b" you want to ensure
    you must specify "a" before "a/b"
    """
    for directory in dirs:
        if not os.path.exists(directory):
            os.makedirs(directory)

def ensure_path(*paths):
    for path in paths:
        d = os.path.dirname(path)
        if len(d)>0 and d!=path:
            ensure_path(d)
        ensure_dir(path)

def move_txt_file(name):
    # split name to [email, version, extension]
    file_info = split_file_name(name)
    if file_info == None:
        return None
##    print file_info
    # create path email/version
    new_dir = os.path.join(file_info[0], file_info[1])
    new_path = os.path.join(new_dir, name)
    ensure_dir(file_info[0], new_dir)
    os.rename(name, new_path)
    return new_dir

def move_archive_file(name, archive):
    """
    extract all the contents of file "name" to email/version/
    move the file "name" to USED_ARCHIVE_DIR
    archive is a class to instantiate the zip/rar file object
    """
    # split name to [email, version, extension]
    file_info = split_file_name(name)
    if file_info == None:
        return None
    new_dir = os.path.join(file_info[0], file_info[1])
    ensure_dir(file_info[0], new_dir)
    zfile = archive(name)
    zfile.extractall(new_dir)
    new_path = os.path.join(USED_ARCHIVE_DIR, name)
##    print 'moving', name, 'to', new_path
    os.rename(name, new_path)
    return new_dir

def move_file(name):
    file_type = FileType.from_file(name)
##    print 'Moving', name, file_type
    if file_type == 'directory':
        #silently ignore
##        print name, 'is a directory, ignore it'
        return None

    if file_type.startswith('ASCII text') or file_type.startswith('Python script') or file_type.startswith('FORTRAN program'):
        return move_txt_file(name)
    elif file_type.startswith('Zip archive data'):
        return move_archive_file(name, zipfile.ZipFile)
    elif file_type.startswith('RAR archive data'):
        return move_archive_file(name, rarfile.RarFile)
    else:
        print name,'has wrong file type:', file_type
        return None

def move_all():
    """
    move all the files in current directory
    """
    moved = []
    notmoved = []
    for f in os.listdir('.'):
        result = move_file(f)
        if result != None:
            moved.append(result)
        else:
            notmoved.append(f)
    return [moved, notmoved]

def purge_osx_junks(root):
    """
    clean the junks generated by Mac OS X filesystem
    """
    for path, dirs, files in os.walk(root):
        if path.startswith('./'):
            path = path[2:]
        if path.find(OSX_JUNK_DIR) >= 0 or path.find(USED_ARCHIVE_DIR)>=0:
            continue
        purged_dir = os.path.join(OSX_JUNK_DIR, path)
        ensured = False
        for f in files+dirs:
            if f in OSX_JUNKS:
                oldpath = os.path.join(path, f)
                newpath = os.path.join(purged_dir, f)
                print oldpath, '->', newpath
                if not ensured:
                    ensure_path(purged_dir)
                    ensured = True
                os.rename(oldpath, newpath)

def flattern_dir(root):
    """
    for any directory, if it only contains a subdir and no files
    flattern it, move all the entries in the only subdir to the parent
    recursively does this
    """
    for path, dirs, files in os.walk(root):
        if path.startswith('./'):
            path = path[2:]
        if path.find(OSX_JUNK_DIR) >= 0 or path.find(USED_ARCHIVE_DIR)>=0:
            continue
        if len(files)==0 and len(dirs)==1:
            subdir = os.path.join(path, dirs[0])
            for f in os.listdir(subdir):
                oldpath = os.path.join(subdir, f)
                newpath = os.path.join(path, f)
                print oldpath, '->', newpath
                os.rename(oldpath, newpath)
            print 'rmdir', subdir
            os.rmdir(subdir)

PROB_NUM_RE = '(?:Problem\s*(?:Set)?|PS|PSet)\s*[1-2]\s*(?:[^a-zA-Z]*Question)?\s*([A-C]|\.[1-3]|\([1-3A-C]\))|Problem\s*([A-C1-3])'
PROB_NUM_MATCHER = re.compile(PROB_NUM_RE, flags=re.IGNORECASE)
def extract_num(comment):
    """
    extract problem number from a one line comment
    case-insensitive
    """

    def translate_num(num):
        if num.startswith('.') or num.startswith('('):
            num = num[1:]
        if num.endswith(')'):
            num = num[:-1]
        num = num.capitalize()
        if num>='A' and num<='C':
            return num
        elif num>='1' and num<='3':
            return chr(ord('A')+ord(num)-ord('1'))
        else:
            print 'Error! Invalid num', num, 'in line', comment
            assert False
            return None

    m = PROB_NUM_MATCHER.search(comment)
    if m != None:
        for num in m.groups():
            if num != None:
                num = translate_num(num)
                return num

    return None

def extract_prob_num(path):
    """
    extract the problem number from a source file "path"
    """
    with open(path) as infile:
        while True:
            line = infile.readline()
            if len(line) == 0:
                return None
            if line.startswith('#'):
                num = extract_num(line)
                if num != None:
                    return num
            else:
                if len(line.strip()) > 0:
                    return None

INFO_FILE_EXT = '.__Prob_Info__'
def get_all_prob_num(root):
    result = []
    no_prob_num = []
    for path, dirs, files in os.walk(root):
        for f in files:
            if f.find(INFO_FILE_EXT) < 0:
                fname = os.path.join(path, f)
                num = extract_prob_num(fname)
                if num == None:
                    print 'Cannot get prob num', fname
                    no_prob_num.append(fname)
                else:
                    result.append(fname)
                    info = fname + INFO_FILE_EXT
                    with open(info, mode='w') as outfile:
                        outfile.write(num)
    return [result, no_prob_num]

PROB_DICT = {'A': '__probA', 'B': '__probB', 'C': '__probC'}
def gather_probs(root, prob_dict):
    """
    prob_dict: {'A': dir_for_probA, 'B': dir_for_probB, ...}
    gather all source code inside root to directories
    """
    for path, dirs, files in os.walk(root):
        for f in files:
            if f.endswith(INFO_FILE_EXT):
                with open(os.path.join(path, f)) as numf:
                    num = numf.read(1)
                sourcefile = f[:-len(INFO_FILE_EXT)]
                oldpath = os.path.join(path, sourcefile)
                if os.path.exists(oldpath):
                    newdir = os.path.join(prob_dict[num], path)
                    ensure_path(newdir)
                    newpath = os.path.join(newdir, sourcefile)
                    print oldpath, '->', newpath
                    os.rename(oldpath, newpath)

TIMEOUT = 5
POLL_INT_START = 0.01
POLL_INT_MAX = 1
def run_py(sourcef, inf, outf):
    if os.path.exists(outf):
        os.remove(outf)
    with open(inf) as infile:
        with open(outf, mode='w') as outfile:
            proc = subprocess.Popen(['python', sourcef], stdin=infile, stdout=outfile, stderr=None)
	    poll_int = POLL_INT_START
            deadline = time.time() + TIMEOUT
            while time.time() < deadline and proc.poll()==None:
                time.sleep(poll_int)
		poll_int = min(POLL_INT_MAX, poll_int*2)
            result = proc.poll()
            if result == None:
                proc.terminate()
    return result
##    cmd = 'timeout ' + str(TIMEOUT) + ' python ' + sourcef + ' < ' + inf + ' > ' + outf
##    print cmd
##    os.system(cmd)

def run_cases(testcase_dir, output_postfix, sourcef):
    for case in os.listdir(testcase_dir):
        infile = os.path.join(testcase_dir, case)
        outfile = sourcef + output_postfix + '.' + case
##        print infile
##        print outfile
##        print sourcef
        run_py(sourcef, infile, outfile)

AVOID_FILENAME = '__'
def test(root, tester):
    for path, dirs, files in os.walk(root):
        for f in files:
            if not AVOID_FILENAME in f:
                sourcef = os.path.join(path, f)
                tester(sourcef)

ENDING_FLOAT_RE = '[-+]?(?:\d*\.\d+|\d+\.?)(?:E[-+]?\d+)?'
ENDING_FLOAT_MATCHER = re.compile(ENDING_FLOAT_RE, flags=re.IGNORECASE)
REMOVE_PROMPT_RE = '[\d]+\s*year'
REMOVE_PROMPT_MATCHER = re.compile(REMOVE_PROMPT_RE, flags=re.IGNORECASE)
def get_ending_float(line):
##    print 'Trying to match', line,
    while True:
        m = REMOVE_PROMPT_MATCHER.search(line)
        if m == None:
            break
        line = line[0:m.start()] + line[m.end()+1:]

    match = None
    for match in ENDING_FLOAT_MATCHER.finditer(line):
        pass
    if match == None:
##        print 'No Match!'
        return None
##    print 'Match:', match.group()
    return float(match.group())

def get_ending_floats(file_path, maxnum=-1):
    """
    read the ending float number from file file_path
    we only read maxnum floats, -1 indicates infinite many
    """
    result = []
    with open(file_path) as f:
        while maxnum != 0:
            line = f.readline()
            if len(line) <= 0:
                break
            fnum = get_ending_float(line)
            if fnum != None:
                --maxnum
                result.append(fnum)
    return result

DEFAULT_OUTPUT_POSTFIX = '.__output__'
DEFAULT_POINT_POSTFIX = '.__point__'
DEFAULT_SCORE_POSTFIX = '.__score__'
class ps1tester:
    """
    Tester for PS1
    """
    def __init__(self, test_case_dir, ans_file_prefix,
                 output_postfix = DEFAULT_OUTPUT_POSTFIX,
                 point_postfix = DEFAULT_POINT_POSTFIX,
                 score_postfix = DEFAULT_SCORE_POSTFIX):
        """
        Initialize the tester
        Read in every answer file and store to a dictionary
            self.cases: infile_name => [infile_path, [floats]]
        """
        self.output_postfix = output_postfix
        self.point_postfix = point_postfix
        self.score_postfix = score_postfix
        self.cases = {}
        self.scores = {}
        for f in os.listdir(test_case_dir):
            infile_path = os.path.join(test_case_dir, f)
            ans_file = ans_file_prefix + f
            self.cases[f] = [infile_path, get_ending_floats(ans_file)]

    def test(self, sourcef):
        score = 0
        summary = ''
        print 'Testing', sourcef
        for key in sorted(self.cases.keys()):
            [inf, ans] = self.cases[key]
            outf = sourcef + self.output_postfix + '.' + key
            result = run_py(sourcef, inf, outf)
            if result == None:
                point = 0
                reason = 'TIMEOUT'
                prompt = 'TIMEOUT'
            elif not os.path.exists(outf):
                point = 0
                reason = 'No Output ' + outf
                prompt = 'NoOut'
            else:
                point = 1
                try:
                    out = get_ending_floats(outf, maxnum=len(ans))
                except:
                    point = 0
                    reason = 'Error reading ' + outf
                    prompt = 'Err'

                if point > 0:
                    prompt = 'ok'
                    reason = ''
                    for i in range(len(ans)):
                        if i >= len(out):
                            point = 0
                            reason = reason + 'None' + '\t #' + str(ans[i]) + ' NA'
                            prompt = 'Shorter'
                        else:
                            reason = reason + str(out[i]) + '\t # ' + str(ans[i])
                            if abs(ans[i] - out[i]) > 0.21:
                                point = 0
                                prompt = 'Wrong'
                                reason = reason + ' Wrong'
                            else:
                                reason = reason + ' OK'
                        reason = reason + '\n'
                    
            score += point
            print prompt
            summary = summary + key + ' : ' + prompt + '\n'
            pointf = sourcef + self.point_postfix + '.' + key
            with open(pointf, mode='w') as pf:
                pf.writelines(str(point) + '\n\nReason:\n' + reason)
        print 'Score:', score
        self.scores[sourcef] = [score, summary]
        scoref = sourcef + self.score_postfix
        with open(scoref, mode='w') as sf:
            sf.writelines(str(score) + '\n\nSummary:\n' + summary)

    def scores(self):
        return self.scores

def test_all(probnum):
    root = PROB_DICT[probnum]
    tester = ps1tester('__input'+probnum, '__answer'+probnum+'/' + probnum + '.ans.')
    test(root, tester.test)
    r = tester.scores
    result_file = os.path.join(root, '__test_result_.txt')
    summary_file = os.path.join(root, '__summary__.txt')
    with open(result_file, mode='w') as rf:
        with open(summary_file, mode='w') as sf:
            for key in sorted(r.keys()):
                [score, summary] = r[key]
                result = key + '\t' + str(score) + '\n'
                rf.writelines(result)
                sf.writelines(result + '\n' + summary + '\n\n')
            
    return tester

def sample_judge(probnum):
    r = test_all(probnum)
    for name, [score, summary] in r.scores.items():
        if score < 7:
            print name, score
    return r

