import time
import base64
import warnings
import numpy as np
#import cv2
from threading import Thread






class AsyncWrite(Thread):
    def __init__(self, png_data, out):
        # calling superclass init
        Thread.__init__(self)
        self.png_data = png_data
        self.out = out

    def run(self):
        f = open(self.out, "wb")
        f.write(self.png_data)
        f.close()

        # waiting for 2 seconds after writing
        # the file
        time.sleep(2)
        print("Finished background file write to",
              self.out)


class Screenshot():
    def __init__(self, driver, screenshots_path):
        self.driver = driver
        self.screenshots = {}
        self.screenshots_sequence = []
        self.screenshots_path = screenshots_path

    def save(self, filename=None):
        if filename == None:
            self.screenshots[filename] = self.driver.get_screenshot_as_base64()
            print "captured"
            base64decoded_image = base64.b64decode(self.screenshots[filename])
            self.screenshots_sequence.append(np.frombuffer(base64decoded_image, dtype='uint8'))
            png_data = base64decoded_image
            background_save = AsyncWrite(png_data, self.screenshots_path + str(time.time()) + ".png")
            background_save.start()
            print "starting save"
            #save_thread = Thread(target=self.save_to_file, args = (str(time.time()) + ".png", png_data, ))
            #save_thread.start()
        else:
            self.screenshots[filename] = self.driver.get_screenshot_as_base64()
            self.screenshots_sequence.append(np.frombuffer(base64.b64decode(self.screenshots[filename]), dtype='uint8'))
            png_data = base64.b64decode(self.screenshots[filename].encode('ascii'))
            self.save_to_file(filename, png_data)

    def save_to_file(self, filename, png_data):
        filename = self.screenshots_path + filename
        if not filename.lower().endswith('.png'):
            warnings.warn("name used for saved screenshot does not match file "
                          "type. It should end with a `.png` extension", UserWarning)
        png = png_data
        try:
            with open(filename, 'wb') as f:
                f.write(png)
        except IOError:
            return False
        finally:
            del png
        return True

    def compare_previous_screenshot(self):
        screenshot_base64 = np.frombuffer(base64.b64decode(self.driver.get_screenshot_as_base64()), dtype='uint8')
        image1 = cv2.imdecode(self.screenshots_sequence[-1], 1)
        image2 = cv2.imdecode(screenshot_base64, 1)
        err = np.sum((image1.astype("float") - image2.astype("float")) ** 2)
        err /= float(image1.shape[0] * image2.shape[1])
        print err, screenshot_base64
        return err, screenshot_base64


    