import base64
import os

import folium
from folium.plugins.heat_map import HeatMap
from folium.utilities import temp_html_filepath


def test_heat_map_with_weights(driver):
    """Verify that HeatMap uses weights in data correctly.

    This test will fail in non-headless mode because window size will be different.

    """
    m = folium.Map((0.5, 0.5), zoom_start=8, tiles=None)
    HeatMap(
        # make four dots with different weights: 1, 1, 1.5 and 2.
        data=[
            (0, 0, 1.5),
            (0, 1, 1),
            (1, 0, 1),
            (1, 1, 2),
        ],
        radius=70,
        blur=50,
    ).add_to(m)
    html = m.get_root().render()
    with temp_html_filepath(html) as filepath:
        driver.get_file(filepath)
        assert driver.wait_until('.folium-map')
        driver.verify_js_logs()
    canvas = driver.wait_until('canvas.leaflet-heatmap-layer')
    assert canvas
    # get the canvas as a PNG base64 string
    canvas_base64 = driver.execute_script(
        "return arguments[0].toDataURL('image/png').substring(21);", canvas)
    screenshot = base64.b64decode(canvas_base64)
    path = os.path.dirname(__file__)
    with open(os.path.join(path, 'test_heat_map_selenium_screenshot.png'), 'rb') as f:
        screenshot_expected = f.read()
    if hash(screenshot) != hash(screenshot_expected):
        print(screenshot)
        assert False, 'screenshot is not as expected'

@@ -14,18 +14,51 @@
 %0Aimport 
-os
+io%0Aimport os%0A%0Afrom PIL import Image
 %0A%0Aimport
@@ -1113,16 +1113,22 @@
 reenshot
+_bytes
  = base6
@@ -1154,16 +1154,74 @@
 base64)%0A
+    screenshot = Image.open(io.BytesIO(screenshot_bytes))%0A
     path
@@ -1370,28 +1370,37 @@
 d = 
-f.read(
+Image.open(f
 )%0A    
+ 
-if hash
+   if list
 (scr
@@ -1410,17 +1410,27 @@
 shot
+.getdata()
 ) != 
-hash
+list
 (scr
@@ -1449,37 +1449,25 @@
 cted
-):%0A        print(screenshot)%0A
+.getdata()):%0A    
     

fix test heat map selenium (#1586)