import dominate
from dominate.tags import meta, h3, table, tr, td, p, a, img, br
import glob
import re
import os
import click
import shutil
from collections import OrderedDict
import numpy as np


def atoi(text):
    return int(text) if text.isdigit() else text

def natural_keys(text):
    '''
    alist.sort(key=natural_keys) sorts in human order
    http://nedbatchelder.com/blog/200712/human_sorting.html
    (See Toothy's implementation in the comments)
    '''
    return [atoi(c) for c in re.split('(\d+)', text)]

def natural_sort(items):
    items.sort(key=natural_keys)


IMG_EXTENSIONS = [
    '.jpg', '.JPG', '.jpeg', '.JPEG',
    '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP', '.tiff', '.webp'
]


def is_image_file(filename):
    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)


class HTML:

    def __init__(self, title, web_dir_prefix='web_'):
        """Initialize the HTML classes

        Parameters:
            web_dir (str) -- a directory that stores the webpage. HTML file will be created at <web_dir>/index.html; images will be saved at <web_dir/images/
            title (str)   -- the webpage name
            reflect (int) -- how often the website refresh itself; if 0; no refreshing
        """
        self.title = title
        self.web_dir = web_dir_prefix

        self.doc = dominate.document(title=title)

    def read_images(self, image_globs):
        image_path_dict = OrderedDict()
        for image_glob in image_globs:
            if not '*' in image_glob and os.path.isdir(image_glob):
                image_glob = image_glob + '/*'
            file_list = glob.glob(os.path.expanduser(image_glob))
            image_list = []
            for file_path in file_list:
                if is_image_file(file_path):
                    image_list.append(file_path)
                else:
                    print('%s does not have valid image extension, so we exclude...' % file_path)
            
            natural_sort(image_list)
            image_path_dict[image_glob] = image_list
        
        num_rows_image_table = len(image_list)
        last_glob = image_glob
        for image_glob, image_list in image_path_dict.items():
            print('Using %s, we found %d images' % (image_glob, len(image_list)))

            if len(image_list) == 0:
                print('%s does not contain any images, fix your command' % image_glob)
                exit(0)

            if len(image_list) != num_rows_image_table:
                print('%d images in %s does not match %d images of %s' %
                      (len(image_list), image_glob, num_rows_image_table, last_glob))
                
                if not click.confirm('Do you still want to continue? (default: no)', default=False):
                    exit(0)

        self.image_path_dict = image_path_dict
        return image_path_dict

    def backup_images(self, image_path_dict=None):
        if image_path_dict is None:
            image_path_dict = self.image_path_dict

        assert image_path_dict is not None
        
        os.makedirs(self.web_dir, exist_ok=True)

        backup_image_path_dict = OrderedDict()
        for image_glob, image_list in image_path_dict.items():            
            subdir_name = os.path.abspath(image_glob.replace("*", "$"))[1:] # remove first /
            subdir_path = os.path.join(self.web_dir, subdir_name)
            os.makedirs(subdir_path, exist_ok=True)

            backup_image_paths = []
            for image_path in image_list:
                backup_image_path = os.path.join(subdir_path, os.path.basename(image_path))
                shutil.copyfile(image_path, backup_image_path)

                # path from index.html's location
                relative_path = os.path.join(subdir_name, os.path.basename(image_path))
                backup_image_paths.append(relative_path)
                
            backup_image_path_dict[image_glob] = backup_image_paths

        self.backup_image_path_dict = backup_image_path_dict
        return backup_image_path_dict


    def add_header(self, text):
        """Insert a header to the HTML file

        Parameters:
            text (str) -- the header text
        """
        with self.doc:
            h3(text)


    def add_image_rows(self, image_path_dict=None, max_image_rows=99999,
                       image_width=256, column_names=['GT'], shuffle=False):
        if image_path_dict is None:
            image_path_dict = self.image_path_dict

        assert image_path_dict is not None

        columns = list(image_path_dict.keys())
        first_key = columns[0]

        num_rows = [len(pathlist) for pathlist in image_path_dict.values()]
        max_rows = np.max(num_rows)        
        max_rows = min(max_rows, max_image_rows)

        for colname, num_row in zip(columns, num_rows):
            print('%s: %d rows' % (colname, num_row))

        print('inserting %d rows...' % max_rows)

        rows = list(range(max_rows))
        if shuffle:
            import random
            random.shuffle(rows)
        
        for i in rows:
            imgs = [image_path_dict[key][i] for key in columns]
            txts = [column_names[0] + ' ' + os.path.basename(image_path_dict[first_key][i])] + \
                   column_names[1:]
            self.add_image_row(imgs, txts, image_width)
        

    def add_image_row(self, ims, txts, width=400):
        """add images to the HTML file

        Parameters:
            ims (str list)   -- a list of image paths
            txts (str list)  -- a list of image names shown on the website
            links (str list) --  a list of hyperref links; when you click an image, it will redirect you to a new page
        """
        self.t = table(border=1, style="table-layout: fixed;")  # Insert a table
        self.doc.add(self.t)
        with self.t:
            with tr():
                for im, txt in zip(ims, txts):
                    with td(style="word-wrap: break-word;", halign="center", valign="top"):
                        with p():
                            with a(href=im):
                                img(style="width:%dpx" % width, src=im)
                            br()
                            p(txt)

    def save(self):
        """save the current content to the HMTL file"""
        html_file = '%s/index.html' % self.web_dir        
        f = open(html_file, 'wt')
        f.write(self.doc.render())
        f.close()
        print(html_file)


    


if __name__ == '__main__':  # we show an example usage here.

    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--web_dir_prefix', type=str, default='web_')
    parser.add_argument('--name', type=str, required=True, help='the name of this image table')
    parser.add_argument('-d', '--dirs', nargs='+', required=True, help='the directories or wildcard expression that will be used by glob.glob()')
    parser.add_argument('-l', '--labels', nargs='+', required=True, help='the column names')
    parser.add_argument('--max_rows', type=int, default=99999)
    parser.add_argument('--image_width', type=int, default=256)
    parser.add_argument('--shuffle', action='store_true')
    opt = parser.parse_args()

    html = HTML(opt.name, opt.web_dir_prefix)
    html.read_images(opt.dirs)
    # html.backup_images()
    html.add_header(opt.name)
    html.add_image_rows(max_image_rows=opt.max_rows, image_width=opt.image_width, column_names=opt.labels, shuffle=opt.shuffle)
    html.save()