import shapefile
import math
from PIL import Image, ImageDraw
"""
Python图像库PIL(Python Image Library)是python的第三方图像处理库，
但是由于其强大的功能与众多的使用人数，几乎已经被认为是python官方图像处理库了。
"""


# 地理坐标转换为屏幕坐标
def world2screen(bbox, w, h, x, y):
    minx, miny, maxx, maxy = bbox
    xdist = maxx - minx
    ydist = maxy - miny
    xratio = w / xdist
    yratio = h / ydist
    px = int(w - ((maxx - x) * xratio))
    py = int(((maxy - y) * yratio))
    return px, py  # 为了适应pngcanvas的坐标系，对y做了翻转


def main():
    r = shapefile.Reader("shapefiles/GIS_CensusTract/GIS_CensusTract_poly")
    width = 600
    height = 400
    # 初始化PIL库的Image对象
    img = Image.new("RGB", (width, height), (255, 255, 255))
    # PIL库的Draw模块用于填充多边形
    draw = ImageDraw.Draw(img)
    # 获取人口和区域
    pop_index = None
    area_index = None

    # 绘制人口普查区域阴影
    for i, f in enumerate(r.fields):
        if f[0] == "POPULAT11":
            pop_index = i-1
        elif f[0] == "AREASQKM":
            area_index = i-1
    # 绘制多边形
    for sr in r.shapeRecords():
        density = sr.record[pop_index]/sr.record[area_index]

        # weight 可以用来配置人口的颜色深浅
        weight = min(math.sqrt(density/80.0), 1.0)*50
        # print(density,"----",weight)
        R = int(205 - weight)
        G = int(215 - weight)
        B = int(245 - weight)
        pixels = []
        for x, y in sr.shape.points:
            (px, py) = world2screen(r.bbox, width, height, x, y)
            pixels.append((px, py))
        draw.polygon(pixels, fill=(R, G, B), outline=(255, 255, 255))
    # img.save("choropleth.png")
    img.show()
    img.close()
    r.close()


if __name__ == '__main__':
    main()