import codecs
import os
from bs4 import BeautifulSoup as bs
from bs4 import NavigableString as ns
import argparse

###### This script is used to generate an HTML report of the SR results ######

parser = argparse.ArgumentParser(description="HTML Report Generator")
parser.add_argument("--sr_results_dir", default='./', help="The path of the directory where the current results are located (assumed to be in .png format") # Mandatory argument
parser.add_argument("--gt_dir",  default = './GT', help="The directory where the ground truth images are located") # Mandatory argument
parser.add_argument("--gt_extension",default = '.jpg', help="The extension of the ground truth images. Default is .jpg ") # Optional argument
###### RUN THIS SCRIPT AFTER RUNNING INFERENCE.PY ######

###### This script requires the following directory format to generate a report. This format is automatically generated if you run the script inference.py for your loss functions. ######


###### SR_Results_Directory -> L1_Loss -> image_1 -> image_1_Bicubic.png
######                                            -> image_1_SR.png
######                                            -> PSNR.txt
######                                    image_2 -> image_2_Bicubic.png
######                                            -> image_2_SR.png
######                                            -> PSNR.txt
######                                      .
######                                      .
######                                      .
######                                      .
######                      -> MSE_Loss -> image_1 -> image_1_Bicubic.png
######                                             -> image_1_SR.png
######                                             -> PSNR.txt
######                                     image_2 -> image_2_Bicubic.png
######                                             -> image_2_SR.png
######                                             -> PSNR.txt
######                                      .
######                                      .
######                                      .
######                                      .
######
######                       .
######                       .
######                       .
######                       .
######
opt = parser.parse_args()
sr_results_dir = opt.sr_results_dir
gt_dir = opt.gt_dir
gt_extension = opt.gt_extension

#sr_results_dir = '/Users/varunbabbar/Desktop/Test' #The path of the directory where the current results are located (assumed to be in .png format)


#gt_dir = '/Users/varunbabbar/Desktop/Summer_Stuff_2020/Report_Code_HTML/SR_Results/GT/' # The directory where the ground truth images are located


#gt_extension = '.jpg' # The extension of the ground truth images


losses = []
for i in os.listdir(sr_results_dir):
    if os.path.isdir(os.path.join(sr_results_dir,i)) and '.ipynb' not in i and 'GT' not in i:
        losses.append(i)
soup = bs()
html = soup.new_tag("html")
body = soup.new_tag("body")
script = soup.new_tag("script")
script.attrs['src'] = "https://ajax.googleapis.com/ajax/libs/jquery/1.9.1/jquery.min.js"
h1 = soup.new_tag("h1")
h1.append("Select the Loss Function results to display! ")
p1 = soup.new_tag("p")
br = soup.new_tag("br")
p1.append("Note: All values of NIMA and PSNR quoted are taken from the epoch where the test loss function was the lowest")
soup.append(html)
html.append(body)
body.append(script)
body.append(h1)
body.append(p1)
body.append(br)
p1 = soup.new_tag("p")
p1.append("Click on any loss function image to view the ground truth image!")
body.append(p1)
#body.append(br)

div_checkbox = soup.new_tag("div")
div_checkbox.attrs['method'] = "get"
div_checkbox.attrs['id'] = "form"
div_checkbox.attrs['value'] = "form"
for loss in losses:
    input_tagx = soup.new_tag("input")
    input_tagx.attrs['class'] = "checkBoxes"
    input_tagx.attrs['type'] = "checkbox"
    input_tagx.attrs['name'] = "Loss Function"
    input_tagx.attrs['id'] = loss
    input_tagx.attrs['value'] = loss
    input_tagx.append(loss)
    div_checkbox.append(input_tagx)
    div_checkbox.append(soup.new_tag("br"))
    
input_tagx = soup.new_tag("input")
input_tagx.attrs['id'] = "#toggle"
input_tagx.attrs['type'] = "submit"
input_tagx.attrs['onclick'] = "return ColumnDelete()"
input_tagx.attrs['value'] = "submit"
div_checkbox.append(input_tagx)

input_tagx = soup.new_tag("input")
input_tagx.attrs['id'] = "select_all"
input_tagx.attrs['type'] = "button"
input_tagx.attrs['name'] = "select_all"
input_tagx.attrs['onclick'] = "return toggle()"
input_tagx.attrs['value'] = "Select All"
div_checkbox.append(input_tagx)
body.append(div_checkbox)
script = soup.new_tag("script")
script.append(""" function toggle() { """)
for loss in losses:
    script.append("""document.getElementById("{}").checked = !document.getElementById("{}").checked""".format(loss,loss))
    
script.append(""" }  $(document).ready(function() {
           $('#select_all').click(function() {
               $('.checkBoxes').each(function() {
                   $(this).attr('checked',!$(this).attr('checked'));
               });
           });
       });""")
body.append(script)
script = soup.new_tag("script")
script.attrs['type'] = "text/javascript"
script.append("""function GetElementInsideContainer(containerID, childID) {
    var elm = document.getElementById(childID);
    var parent = elm ? elm.parentNode : {};
    return (parent.id && parent.id === containerID) ? elm : {};
}

function testWhite(x) {
    var white = new RegExp(/^\s$/);
    return white.test(x.charAt(0));
};""")

script.append("""function getOffset( el ) {
    var _x = 0;
    var _y = 0;
    while( el && !isNaN( el.offsetLeft ) && !isNaN( el.offsetTop ) ) {
        _x += el.offsetLeft - el.scrollLeft;
        _y += el.offsetTop - el.scrollTop;
        el = el.offsetParent;
    }
    return { top: _y, left: _x };
}



var para = document.createElement("h3");


function ModifypecificColumns(table_id, loss,remove) {
    elem = document.getElementById(table_id);
    for (i = 0; i < elem.childElementCount; i++) {

      var child = elem.children[i]
      for(j= 0; j < child.childElementCount; j++){
          var grandchild = child.children[j]
          if(grandchild.id == loss){
              if (remove){
                  grandchild.style.display="none"
              }
              else{
                  grandchild.style.display=""
              }
              
          }
      }
                
      }
}

function AddPSNRtoTable(table_id,loss){
    elem = document.getElementById(table_id);
       for (i = 0; i < elem.childElementCount; i++) {

         var child = elem.children[i]
         
         for(j= 0; j < child.childElementCount; j++){
             var grandchild = child.children[j]
             
             if(grandchild.id == loss && child.id != "Heading"){
                 
                 for(k =0;k <grandchild.childElementCount; k++){
                     ggchild = grandchild.children[k]
                     if(ggchild.id == loss || ggchild.id == loss ){
                         
                         if (loss == "Ground Truth" || loss == "Bicubic"){
                             img = child.id
                             paragraph = ggchild.getElementsByTagName("P")[0]
                             
                         }
                         else{
                         try{
                             img = child.id
                             //readTextFile(loss+"/" + img + "/" + "PSNR.txt",img+loss)
                            // psnr_element = document.getElementById(img+loss)
                             
                             //psnr = psnr_element.innerText
                             psnr = ""
                             paragraph = ggchild.getElementsByTagName("P")[0]
                             //paragraph.innerText = psnr
                             var nima_loc = psnr.indexOf("NIMA:");
                             var nima = psnr.substring(nima_loc+5,nima_loc + 11)
                             var psnr_loc = psnr.indexOf("PSNR:");
                             var psnr = psnr.substring(psnr_loc+5,psnr_loc + 11)
                             paragraph.innerText = "PSNR: " + psnr + "   " +  "NIMA: " + nima
                         }
                         catch{
                             paragraph = ggchild.getElementsByTagName("P")[0]
                             //paragraph.innerText = ""
                             
                         }
                         }
                         
                     }
                 }
             }
         }
                   
         }
}""")

body.append(script)
style = soup.new_tag("style")
style.attrs['type'] = "text/css"
style.append(""".Table
{
    display: table;
    overflow-x: scroll;
    white-space: nowrap;
    resize:both;
}
.Title
{
    display: table-caption;
    text-align: center;
    font-weight: bold;
    font-size: larger;
    text-align: center;
}
.Heading
{
    display: table-row;
    font-weight: bold;
    text-align: center;
}
.Row
{
    display: table-row;
    text-align: center;
}
.Cell
{
    display: table-cell;
    border: solid;
    border-width: thin;
    padding-left: 5px;
    padding-right: 5px;
    text-align: center;
    vertical-align: middle;
}
img {
    max-width:550px;
    max-height:550px;
    width: auto;
    height: auto;
    overflow: hidden;
}""")

body.append(style)


table_div = soup.new_tag("div")
table_div.attrs['class'] = "Table"
table_div.attrs['id'] = "Table"
folder_sampler_dir = os.path.join(sr_results_dir,losses[0])
images = [i for i in os.listdir(folder_sampler_dir) if os.path.isdir(os.path.join(folder_sampler_dir,i)) and '.ipynb' not in i]
### This needs to be done as many times as there are images in the HTML file
plot_metrics = ['LOss','PSNR','NIMA']

for image in images:
    heading_div = soup.new_tag("div")
    heading_div.attrs['class'] = "Heading"
    heading_div.attrs['id'] = "Heading"
    div_list = []
    for loss in losses:
        cell_div = soup.new_tag("div")
        cell_div.attrs['class'] = "Cell"
        cell_div.attrs['id'] = loss
        loss = loss.replace("_"," ")
        cell_div.append(loss)
        div_list.append(cell_div)
    
    cell_div = soup.new_tag("div")
    cell_div.attrs['class'] = "Cell"
    cell_div.attrs['id'] = "Ground Truth"
    cell_div.append("Ground Truth")
    div_list.insert(0, cell_div)

    
    for div in div_list:
        heading_div.append(div)
    table_div.append(heading_div)
    row_div = soup.new_tag("div")
    row_div.attrs['class'] = "Row"
    row_div.attrs['id'] = image
    div_list = []
    for loss in losses:
        cell_div = soup.new_tag("div")
        cell_div.attrs['class'] = "Cell"
        cell_div.attrs['id'] = loss
        image_div = soup.new_tag("div")
        image_div.attrs['class'] = "image"
        image_div.attrs['id'] = loss
        img_div = soup.new_tag("img")
        img_div.attrs['src'] = os.path.join(loss,image,image+".png")
        cell_div.append(soup.new_tag("br"))
        image_div.append(img_div)
        cell_div.append(image_div)
        div_list.append(cell_div)
    
    cell_div = soup.new_tag("div")
    cell_div.attrs['class'] = "Cell"
    cell_div.attrs['id'] = "Ground Truth"
    image_div = soup.new_tag("div")
    image_div.attrs['class'] = "image"
    img_div = soup.new_tag("img")
    img_div.attrs['src'] = os.path.join(gt_dir,image+gt_extension)
    image_div.append(img_div)
    cell_div.append(image_div)
    div_list.insert(0, cell_div)
    for div in div_list:
        row_div.append(div)
    table_div.append(row_div)

body.append(table_div)

script = soup.new_tag("script")
script.append(""" $('img').on({{
    'mousedown': function(event) {{
        switch (event.which) {{
            case 1:

                var parent_id = $(this).parent().attr("id")
                var para = $(this).parent().parent().children()[0];
                var image_number = $(this).parent().parent().parent().attr("id");
                console.log(image_number)
                var image = $(this).attr('src')
                if (parent_id == void 0){{
                    parent_id = 'Ground Truth'
                }}
                
                if (parent_id!='Ground Truth' && parent_id != 'Bicubic' && !image.includes('PSNR') &&!image.includes('LOss')  &&  !image.includes('NIMA')){{
                    console.log(parent_id)
                var slash_index = image.indexOf(parent_id+"/");
                
                var image_value_string = image.substring(slash_index+(parent_id+"/").length,image.length);
                
                if (slash_index == -1){{
                    var slash_index = image.indexOf("{}");
                    var image_value_string = image.substring(slash_index+("{}").length,image.length)
                }}
                else{{
                    para.innerText = 'Ground Truth Image Shown'
                }}
                
                if (image_value_string.substring(image_value_string.length-3,image_value_string.length) == 'jpg' || image_value_string.substring(image_value_string.length-3,image_value_string.length) == 'png'){{
                    para.innerText = 'Current Loss Shown: ' + parent_id
                }}
                var src = ($(this).attr('src') === parent_id+'/'+image_number+'/'+image_number+'.png')
                     ? "{}" +  image_number +"{}"
                     : parent_id+'/'+image_number+'/'+image_number+'.png';
                 $(this).attr('src', src);
                 
             }}
    }}
    }}
}});

$('img').on({{
    'mouseup': function(event) {{
        switch (event.which) {{
            case 1:
            var parent_id = $(this).parent().attr("id")
            var para = $(this).parent().parent().children()[0];
            var image_number = $(this).parent().parent().parent().attr("id");
            var image = $(this).attr('src')
            if (parent_id == void 0){{
                parent_id = 'Ground Truth'
            }}
            
            if (parent_id!='Ground Truth' && parent_id != 'Bicubic' && !image.includes('PSNR') &&!image.includes('LOss')  &&  !image.includes('NIMA')){{
            
            var slash_index = image.indexOf(parent_id+"/");
            var image_value_string = image.substring(slash_index+(parent_id+"/").length,image.length);
            console.log(image_value_string)
            if (slash_index == -1){{
                
                var slash_index = image.indexOf("{}");
                var image_value_string = image.substring(slash_index+("{}").length,image.length)
            }}
            else{{
                para.innerText = 'Ground Truth Image Shown'
            }}
           
            if (image_value_string.substring(image_value_string.length-3,image_value_string.length) == 'jpg' || image_value_string.substring(image_value_string.length-3,image_value_string.length) == 'png'){{
                para.innerText = 'Current Loss Shown: ' + parent_id
            }}
            var src = ($(this).attr('src') === parent_id+'/'+image_number+'/'+image_number+'.png')
                 ? "{}" +  image_number +"{}"
                 : parent_id+'/'+image_number+'/'+image_number+'.png';
             $(this).attr('src', src);
             
         }}
    }}
    }}
}});""".format(os.path.join(gt_dir,""),os.path.join(gt_dir,""),os.path.join(gt_dir,""),gt_extension,os.path.join(gt_dir,""),os.path.join(gt_dir,""),os.path.join(gt_dir,""),gt_extension))
script.append("""AddPSNRtoTable("Table","Ground Truth") \n""")
for loss in losses:
    script.append("""AddPSNRtoTable("Table","{}") \n""".format(loss))
script.append("""   """)
script.append(""" function ColumnDelete(){ \n""")
for loss in losses:
    script.append("""var {} = document.getElementById("{}").checked \n""".format(loss,loss))
    script.append(""" if ({}) {{
        ModifypecificColumns("Table","{}",false)
        }}
        else{{
            ModifypecificColumns("Table","{}",true)
        }} \n """.format(loss,loss,loss)
    )
script.append("""} """)
    
body.append(script)

f = open(os.path.join(sr_results_dir,'test_new.html'),'w')
f.write(str(soup.prettify()))
 
#print(soup.prettify())