#!/usr/local/bin/py
import argparse
import hashlib
import logging
import os
import oss2  # pip install oss2
import re

logging.getLogger("requests").setLevel(logging.WARNING)
logging.getLogger('oss2').setLevel(logging.WARNING)
logging.basicConfig(
    level=logging.INFO,
    format='[%(asctime)s][%(levelname)s] %(message)s',
    # filename='/tmp/oss-sync.log'
)

_CACHE = {}
ROOT_API_KEY = os.path.join(os.getenv('HOME'), '.aliyun')

# Doc: https://help.aliyun.com/knowledge_detail/5974206.htm
# 青岛节点外网地址： oss-cn-qingdao.aliyuncs.com
# 青岛节点内网地址： oss-cn-qingdao-internal.aliyuncs.com
#
# 北京节点外网地址：oss-cn-beijing.aliyuncs.com
# 北京节点内网地址：oss-cn-beijing-internal.aliyuncs.com
#
# 杭州节点外网地址： oss-cn-hangzhou.aliyuncs.com
# 杭州节点内网地址： oss-cn-hangzhou-internal.aliyuncs.com
#
# 上海节点外网地址： oss-cn-shanghai.aliyuncs.com
# 上海节点内网地址： oss-cn-shanghai-internal.aliyuncs.com
#
# 香港节点外网地址： oss-cn-hongkong.aliyuncs.com
# 香港节点内网地址： oss-cn-hongkong-internal.aliyuncs.com
#
# 深圳节点外网地址： oss-cn-shenzhen.aliyuncs.com
# 深圳节点内网地址： oss-cn-shenzhen-internal.aliyuncs.com
#
# 美国节点外网地址： oss-us-west-1.aliyuncs.com
# 美国节点内网地址：  oss-us-west-1-internal.aliyuncs.com
#
# 新加坡节点外网地址： oss-ap-southeast-1.aliyuncs.com
# 新加坡节点内网地址：  oss-ap-southeast-1-internal.aliyuncs.com
#
# 原地址oss.aliyuncs.com 默认指向杭州节点外网地址。
# 原内网地址oss-internal.aliyuncs.com 默认指向杭州节点内网地址
API_URL = 'oss-cn-hangzhou.aliyuncs.com'

IGNORE_FILES = (
    '\/\..*$',
    '\.pyc$',
)

def is_in_ignore_files(file_path):
    for p in IGNORE_FILES:
        if re.search(p, file_path):
            return True
    return False

def get_file_md5(file_path):
    hasher = hashlib.md5()
    with open(file_path, 'rb') as f:
        buf = f.read(65536)
        while len(buf) > 0:
            hasher.update(buf)
            buf = f.read(65536)
    return hasher.hexdigest()

def sizeof_fmt(num):
    if num <= 1024:
        return '1 KB'
    for x in ['bytes', 'KB', 'MB', 'GB', 'TB']:
        if num < 1024.0:
            return "%3.1f %s" % (num, x)
        num /= 1024.0

def get_bucket(args):
    if 'bucket' in _CACHE:
        return _CACHE['bucket']

    api_key = open(os.path.join(ROOT_API_KEY, 'apikey')).read().strip()
    api_secret = open(os.path.join(ROOT_API_KEY, 'secretkey')).read().strip()
    auth = oss2.Auth(api_key, api_secret)
    bucket = oss2.Bucket(auth, API_URL, args.bucket)
    _CACHE['bucket'] = bucket
    return bucket

def get_local_objects(target_path):
    objects = {}
    oss_dir = os.path.dirname(__file__)
    if target_path:
        oss_dir = os.path.join(oss_dir, target_path)
    else:
        oss_dir = os.path.join(oss_dir, '.')
    if not os.path.exists(oss_dir):
        return objects

    file_count = 0
    if os.path.isdir(oss_dir):
        for root, dirs, files in os.walk(oss_dir):
            for f in files:
                root = re.sub(r'^\./?', '', root)
                local_path = os.path.join(root, f)
                if is_in_ignore_files(local_path):
                    logging.info('ignored file: {}'.format(local_path))
                    continue
                md5 = get_file_md5(local_path)
                objects[local_path] = md5.upper()
                file_count += 1
    else:
        md5 = get_file_md5(oss_dir)
        local_path = re.sub(r'^\./', '', target_path)
        objects[local_path] = md5.upper()
        file_count += 1
    logging.info('local files: {}'.format(file_count))
    return objects

def get_remote_objects(args):
    objects = {'files': {}, 'etags': {}, 'meta': {}}
    bucket = get_bucket(args)
    marker = None
    file_count = 0
    prefix = re.sub(r'^\./?', '', args.target_path or '')
    while True:
        result = bucket.list_objects(prefix=prefix, max_keys=100, marker=marker)
        for obj in result.object_list:
            if obj.key.endswith('/'):
                continue
            if args.min_size and obj.size < args.min_size:
                continue
            if args.max_size and obj.size > args.max_size:
                continue
            if args.re and not re.search(args.re, obj.key):
                continue
            objects['files'][obj.key] = obj.etag
            objects['etags'][obj.etag] = obj.key
            objects['meta'][obj.key] = obj
            file_count += 1
        marker = result.next_marker
        if not result.is_truncated:
            break
    logging.info('remote files: {}'.format(file_count))
    return objects

def upload_file(local_path, args):
    bucket = get_bucket(args)
    key = re.sub(r'^\./?', '', local_path)
    res = bucket.put_object_from_file(key, local_path)
    if res.status != 200:
        logging.error('Upload {} failed. Exit.'.format(local_path))
        exit(1)

def upload_files_to_oss(args):
    target_path = re.sub(r'^\./?', '', args.target_path)
    logging.info('Uploading/Updating for: {}'.format(target_path))
    los = get_local_objects(target_path)
    if args.check_duplicated:
        ros = get_remote_objects(args)
    else:
        ros = get_remote_objects(args)

    files_need_to_update = []
    files_need_to_upload = []

    for local_path in los.keys():
        md5 = los[local_path]
        if md5 in ros['etags']:
            logging.info('* Identical file found:')
            logging.info('* @ {}'.format(ros['etags'][md5]))
            continue

        if local_path not in ros['files']:
            size = sizeof_fmt(os.path.getsize(local_path))
            files_need_to_upload.append((local_path, size))
        elif ros['files'][local_path] != md5:
            size = sizeof_fmt(os.path.getsize(local_path))
            files_need_to_update.append((local_path, size))

    files_need_to_update.sort()
    files_need_to_upload.sort()

    index = 1
    count = len(files_need_to_update)
    for local_path, size in files_need_to_update:
        if args.no:
            break
        elif args.yes:
            upload_file(local_path, args)
            index += 1
        else:
            print('Q: Do you want to update {}:'.format(local_path))
            response = input()
            while response.lower().strip() not in ('yes', 'no'):
                print('Q: Do you want to update {}:'.format(local_path))
                response = input()
            if response == 'no':
                logging.info('skipped {} by user'.format(local_path))
                continue
            logging.info('= [{}/{}] Updating old file: {} ({})'.format(
                index, count, local_path, size))
            upload_file(local_path, args)
            index += 1

    index = 1
    count = len(files_need_to_upload)
    for local_path, size in files_need_to_upload:
        try:
            logging.info('+ [{}/{}] Uploading new file: {} ({})'.format(
                index, count, local_path, size))
        except:
            pass
        upload_file(local_path, args)
        index += 1

    logging.info('Uploading/Updating Done\n')

def _get_dir_of_file(f):
    return '/'.join(f.split('/')[:-1])

def download_file(oss_path, local_path, args):
    dir_ = _get_dir_of_file(local_path)
    if not os.path.exists(dir_):
        os.makedirs(dir_)
    logging.info('+ Downloading {}'.format(oss_path))
    bucket = get_bucket(args)
    local_path = local_path.encode('utf-8')
    res = bucket.get_object_to_file(oss_path, local_path)
    if res.status != 200:
        logging.error('Download {} failed. Exit.'.format(oss_path))
        exit(1)

def list_files_on_oss(args):
    files = get_remote_objects(args)
    size_total = 0
    for o in files['meta']:
        size_total += files['meta'][o].size
        if args.verbose:
            print('\n- file: {}'.format(o))
            print('- size: {}'.format(sizeof_fmt(files['meta'][o].size)))
            print('- md5: {}'.format(files['meta'][o].etag))

    if not args.verbose:
        keys_to_list = list(files['files'].keys())
        keys_to_list.sort()
        print('== First 3 files:')
        for x in keys_to_list[:3]:
            print('   - {}'.format(x))
        print('== Last 3 files:')
        for x in keys_to_list[-3:]:
            print('   - {}'.format(x))

    print('\n== Total file count: {}'.format(len(files['files'])))
    print('== Total size: {}'.format(sizeof_fmt(size_total)))

def delete_files_from_oss(args):
    files = get_remote_objects(args)
    keys_to_delete = list(files['files'].keys())
    keys_to_delete.sort()
    print('== Will delete {} files:'.format(len(keys_to_delete)))
    print('== First 3 files:')
    for x in keys_to_delete[:3]:
        print('   - {}'.format(x))
    print('== Last 3 files:')
    for x in keys_to_delete[-3:]:
        print('   - {}'.format(x))

    answer = input('== Please enter YES to delete them ALL: ')
    if answer.strip() != 'YES':
        print('\nAction Canceled. Files are safe. Bye.')
        return

    bucket = get_bucket(args)
    count = 0
    for x in keys_to_delete:
        bucket.delete_object(x)
        count += 1
        print('- deleted: {}'.format(x))
    print('\nDeleted {} files.'.format(count))

def download_files_from_oss(args):
    target_path = args.target_path
    if target_path.startswith('./'):
        target_path = target_path[2:]
    if target_path.startswith('/'):
        raise ValueError('Must use relative path')

    oss_dir = os.path.dirname(__file__)
    oss_dir = os.path.join(oss_dir, '.')
    logging.info('Downloading file from: {}'.format(target_path))
    los = get_local_objects(target_path)
    ros = get_remote_objects(args)
    target_files = []
    for obj_key in ros['files']:
        if obj_key in los and ros['files'][obj_key] == los[obj_key]:
            logging.info('= {} exists'.format(obj_key))
            continue
        target_files.append(obj_key)

    target_files.sort()
    for oss_path in target_files:
        local_path = os.path.join(oss_dir, oss_path)
        download_file(oss_path, local_path, args)
    logging.info('Downloading Done\n')

def main():
    parser = argparse.ArgumentParser(description='Use Aliyun-OSS as Dropbox')
    parser.add_argument(
        '--target-path',
        '-p',
        action='store',
        const=None,
        default=None,
        help='Target path to sync/delete files'
    )
    parser.add_argument(
        '--download',
        '-d',
        action='store_true',
        default=False,
        help='Download files from OSS'
    )
    parser.add_argument(
        '--yes',
        action='store_true',
        default=False,
        help='overwrite existing files'
    )
    parser.add_argument(
        '--no',
        action='store_true',
        default=False,
        help='Do NOT overwrite existing files'
    )
    parser.add_argument(
        '--upload',
        '-u',
        action='store_true',
        default=False,
        help='Upload files to OSS'
    )
    parser.add_argument(
        '--listing',
        '-L',
        action='store_true',
        default=False,
        help='List files meta info on OSS'
    )
    parser.add_argument(
        '--min-size',
        type=int,
        default=0,
        help='[Listing] do not list size smaller than this'
    )
    parser.add_argument(
        '--max-size',
        type=int,
        default=0,
        help='[Listing] do not list size bigger than this'
    )
    parser.add_argument(
        '--re',
        type=str,
        default='',
        help='[Listing] filter file name by RE string'
    )
    parser.add_argument(
        '--check-duplicated',
        '-c',
        action='store_false',
        default=True,
        help='Do not upload files already in bucket other dirs'
    )
    parser.add_argument(
        '--bucket',
        '-b',
        required=True,
        help='bucket name to store data',
    )
    parser.add_argument(
        '--delete',
        action='store_true',
        help='To delete files with prefix from OSS',
    )
    parser.add_argument(
        '--verbose',
        '-v',
        action='store_true',
        help='Print more info',
    )

    args = parser.parse_args()
    if args.listing:
        list_files_on_oss(args)
    elif args.download:
        download_files_from_oss(args)
    elif args.delete:
        delete_files_from_oss(args)
    else:
        upload_files_to_oss(args)

if __name__ == "__main__":
    main()


"""Computes eigenvalues and eigenvectors of the PMI similarity matrices for a given attribute type. Saves the results of this along with kMeans clustering of the attributes, and the assignment of graph nodes to clusters."""

import pickle
import time
import numpy as np
import pandas as pd
import optparse
from scipy.sparse import coo_matrix, diags
from sklearn.cluster import KMeans
from gplus import *

def generate_cluster_report(attr_analyzer, attr_type, cluster_labels, topN = 30):
    """Given the AttributeAnalyzer, attr_type, and a list of cluster labels (corresponding to the attribute vocab indices only), generates a report listing the top N members of each cluster, and the frequency and prevalence (relative frequency) of each attribute in the data set. Orders the clusters by total occurrences of attributes in each cluster. If topN = None, list all the attributes in each cluster."""
    attr_freq_dict = attr_analyzer.attr_freqs_by_type[attr_type]
    total_attr_freqs = sum(attr_freq_dict.values())
    pfa = attr_analyzer.pairwise_freq_analyzers[attr_type]
    attr_indices, attr_vocab = get_attr_indices(pfa, attr_analyzer.attributed_nodes)
    unique_cluster_labels = set(cluster_labels)
    # compute vocab lists for each cluster
    attr_vocab_by_cluster = dict((lab, []) for lab in unique_cluster_labels)
    for (i, lab) in enumerate(cluster_labels):
        v = attr_vocab[i]
        if v.startswith('*???*'):
            continue
        freq = attr_freq_dict[v]
        attr_vocab_by_cluster[lab].append((v, freq, freq / total_attr_freqs))
    # sort vocab lists by decreasing frequencies
    for lab in unique_cluster_labels:
        attr_vocab_by_cluster[lab].sort(key = lambda item : item[1], reverse = True)
    # total number of occurrences of any attribute in each cluster
    total_freqs_by_cluster = dict((lab, sum([item[1] for item in attr_vocab_by_cluster[lab]])) for lab in unique_cluster_labels)
    info_by_cluster = dict((lab, dict()) for lab in unique_cluster_labels)
    # create a DataFrame for each cluster listing the top N vocab items in order with their frequencies and prevalences
    for lab in unique_cluster_labels:
        df = pd.DataFrame(attr_vocab_by_cluster[lab], columns = ['attribute', 'frequency', 'prevalence'])
        info_by_cluster[lab]['df'] = df if (topN is None) else df[:topN]
        info_by_cluster[lab]['size'] = len(attr_vocab_by_cluster[lab])
        info_by_cluster[lab]['totalFreq'] = total_freqs_by_cluster[lab]
        info_by_cluster[lab]['totalPrevalence'] = sum(df['prevalence'])
    # sort clusters by decreasing number of occurrences
    sorted_clusters_with_total_freqs = sorted(total_freqs_by_cluster.items(), key = lambda item : item[1], reverse = True)
    # generate report
    num_attrs = len(attr_vocab)
    s = ''
    for (lab, freq) in sorted_clusters_with_total_freqs:
        info = info_by_cluster[lab]
        width = 12 + len(str(lab))
        s += '#' * width + '\n'
        s += '# ' + 'CLUSTER ' + str(lab) + ' #\n'
        s += '#' * width + '\n\n'
        s += 'attribute prevalence  = %6d / %6d = %f\n' % (info['size'], num_attrs, info['size'] / num_attrs)
        s += 'occurrence prevalence = %6d / %6d = %f\n\n' % (info['totalFreq'], total_attr_freqs, info['totalPrevalence'])
        s += info['df'].to_string(index = False) + '\n\n\n'
    return s

# save off:
# matrix or LinearOperator for similarity matrix
# eigenvalues and scree plot
# embedded vectors corresponding to attributes
# kmeans clusters corresponding to attributes
# report of top clusters
# mappings from nodes to clusters of the given attribute type
def main():
    p = optparse.OptionParser()
    p.add_option('--attr_type', '-a', type = str, help = 'attribute type')
    p.add_option('-p', type = str, help = 'PMI type (PMIs, NPMI1s, or NPMI2s)')
    p.add_option('-e', type = str, help = 'embedding (adj, normlap, regnormlap)')
    p.add_option('-s', action = 'store_true', default = False, help = 'normalize in sphere')
    p.add_option('-d', type = float, help = 'smoothing parameter')
    p.add_option('-k', type = int, help = 'number of eigenvalues')
    p.add_option('-c', type = int, help = 'number of kmeans clusters')
    p.add_option('-t', type = float, default = None, help = 'tolerance for eigsh')
    p.add_option('-v', action = 'store_true', default = False, help = 'save scree plot')
    opts, args = p.parse_args()

    attr_type = opts.attr_type
    sim = opts.p
    embedding = opts.e
    assert (embedding in ['adj', 'normlap', 'regnormlap'])
    sphere = opts.s
    delta = opts.d
    k = opts.k
    nclusts = opts.c
    tol = opts.t
    save_plot = opts.v
    topN = 50  # for the report
    assert (((sim == 'PMIs') or (delta == 0)) and (sim in ['PMIs', 'NPMI1s', 'NPMI2s']))

    data_folder = 'gplus0_lcc/data/PMI/'
    report_folder = 'gplus0_lcc/reports/PMI/'
    plot_folder = 'gplus0_lcc/plots/PMI/'
    file_prefix1 = ('%s_%s_%s_delta' % (attr_type, sim, embedding)) + str(delta) + ('_k%d' % k)
    file_prefix2 = ('%s_%s_%s_delta' % (attr_type, sim, embedding)) + str(delta) + ('_k%d%s_c%d' % (k, '_normalized' if sphere else '', nclusts))

    print_flush("\nLoading AttributeAnalyzer...")
    a = AttributeAnalyzer()
    a.load_pairwise_freq_analyzer(attr_type)
    a.make_attrs_by_node_by_type()
    attrs_by_node = a.attrs_by_node_by_type[attr_type]
    pfa = a.pairwise_freq_analyzers[attr_type]
    n = pfa.num_vocab
    tol = (1.0 / n) if (tol is None) else tol  # use 1/n instead of machine precision as default tolerance
    attr_indices, attr_vocab = get_attr_indices(pfa, a.attributed_nodes)

    try:
        print_flush("\nLoading labels from '%s%s_labels.csv'..." % (data_folder, file_prefix2))
        labels = np.loadtxt('%s%s_labels.csv' % (data_folder, file_prefix2), dtype = int)
        print_flush("\nLoading cluster centers from '%s%s_cluster_centers.csv'..." % (data_folder, file_prefix2))
        cluster_centers = np.loadtxt('%s%s_cluster_centers.csv' % (data_folder, file_prefix2), delimiter = ',')
        print_flush("\nLoading eigenvalues from '%s%s_eigvals.csv'..." % (data_folder, file_prefix1))
        eigvals = np.loadtxt('%s%s_eigvals.csv' % (data_folder, file_prefix1), delimiter = ',')
        print_flush("\nLoading embedded features from '%s%s_features.pickle'..." % (data_folder, file_prefix1))
        features = pickle.load(open('%s%s_features.pickle' % (data_folder, file_prefix1), 'rb'))
        if sphere:
            for i in range(len(attr_indices)):
                features[i] = normalize(features[i])
    except FileNotFoundError:
        print_flush("Failed to load.")
        try:
            print_flush("\nLoading eigenvalues from '%s%s_eigvals.csv'..." % (data_folder, file_prefix1))
            eigvals = np.loadtxt('%s%s_eigvals.csv' % (data_folder, file_prefix1), delimiter = ',')
            print_flush("\nLoading embedded features from '%s%s_features.pickle'..." % (data_folder, file_prefix1))
            features = pickle.load(open('%s%s_features.pickle' % (data_folder, file_prefix1), 'rb'))
        except FileNotFoundError:
            print_flush("Failed to load.")
            print_flush("\nComputing similarity matrix (%s)..." % sim)
            sim_op = pfa.to_sparse_PMI_operator(sim, delta)
            matrix_type = 'adjacency' if (embedding == 'adj') else ('normalized Laplacian' if (embedding == 'normlap') else 'regularized normalized Laplacian')
            print_flush("\nComputing eigenvectors of %s matrix (k = %d)..." % (matrix_type, k))
            if (embedding == 'adj'):
                (eigvals, features) = timeit(eigsh)(sim_op, k = k, tol = tol)
                features = np.sqrt(np.abs(eigvals)) * features  # scale the feature columns by the sqrt of the eigenvalues
            elif (embedding == 'normlap'):
                normlap = SparseNormalizedLaplacian(sim_op)
                (eigvals, features) = timeit(eigsh)(normlap, k = k, tol = tol)
            elif (embedding == 'regnormlap'):
                regnormlap = SparseRegularizedNormalizedLaplacian(sim_op)
                (eigvals, features) = timeit(eigsh)(regnormlap, k = k, tol = tol)
            features = features[attr_indices, :]  # free up memory by deleting embeddings of nodes with no attributes
            np.savetxt('%s%s_eigvals.csv' % (data_folder, file_prefix1), eigvals, delimiter = ',')
            pickle.dump(features, open('%s%s_features.pickle' % (data_folder, file_prefix1), 'wb'))
        if sphere:  # normalize the features to have unit norm (better for kMeans)
            for i in range(len(attr_indices)):
                features[i] = normalize(features[i])
        km = KMeans(nclusts)
        print_flush("\nClustering attribute feature vectors into %d clusters using kMeans..." % nclusts)
        labels = timeit(km.fit_predict)(features)
        # save the cluster labels
        np.savetxt('%s%s_labels.csv' % (data_folder, file_prefix2), np.array(labels, dtype = int), delimiter = ',', fmt = '%d')
        # save the cluster centers
        cluster_centers = km.cluster_centers_
        np.savetxt('%s%s_cluster_centers.csv' % (data_folder, file_prefix2), cluster_centers, delimiter = ',')
        # save the attribute cluster report
        with open('%s%s_cluster_report.txt' % (report_folder, file_prefix2), 'w') as f:
            f.write(generate_cluster_report(a, attr_type, labels, topN))

    if save_plot:
        print_flush("\nSaving scree plot to '%s%s_screeplot.png'..." % (plot_folder, file_prefix1))
        scree_plot(eigvals, show = False, filename = '%s%s_screeplot.png' % (plot_folder, file_prefix1))

    print_flush("\nAssigning cluster labels to each node...")
    indices_by_vocab = dict((v, i) for (i, v) in enumerate(attr_vocab))
    centers = [normalize(center) for center in cluster_centers] if sphere else cluster_centers
    def assign_cluster(node):
        """Assigns -1 to a node with no attribute present. Otherwise, takes the cluster whose center is closest to the mean of the attribute vectors. Uses cosine distance if sphere = True, otherwise Euclidean distance."""
        if (node not in attrs_by_node):
            return -1
        else:
            attrs = list(attrs_by_node[node])
            if (len(attrs) == 1):
                return labels[indices_by_vocab[attrs[0]]]
            else:
                vec = np.zeros(k, dtype = float)
                for attr in attrs:
                    vec += features[indices_by_vocab[attr]]
                vec /= len(attrs)
                if sphere:
                    vec = normalize(vec)
                    sims = [np.dot(vec, center) for center in centers]
                else:
                    sims = [-np.linalg.norm(vec - center) for center in centers]
                max_index, max_sim = -1, -float('inf')
                for (i, sim) in enumerate(sims):
                    if (sim > max_sim):
                        max_index = i
                        max_sim = sim
                return max_index

    # save file with the list of cluster labels for each node
    clusters_by_node = [assign_cluster(i) for i in range(a.num_vertices)]
    np.savetxt('%s%s_node_labels.csv' % (data_folder, file_prefix2), np.array(clusters_by_node, dtype = int), delimiter = ',', fmt = '%d')
    print_flush("\nDone!")

if __name__ == "__main__":
    main()