from owslib.wms import WebMapService

# Initialisiere Web Mapping Service fuer das Land Bayern
wms = WebMapService('https://geoservices.bayern.de/od/wms/dop/v1/dop20?')from owslib.wms import WebMapService

# Definiere eine Funktion, welche RGBI-Daten herunterlaedt und speichert, basierend auf einer bounding box (Gauss-Kruger zone 4 Koordinaten)
def get_wms_tile(bbox, 
                 gsd=0.2,
                 url='https://geoservices.bayern.de/od/wms/dop/v1/dop20?', 
                 srs='EPSG:31468', 
                 format='image/tiff',
                 targetdir='./'):
    
    # open wms
    wms = WebMapService(url)

    # determine image shape in pixels
    width = int((bbox[2]-bbox[0])/0.2)
    height = int((bbox[3]-bbox[1])/0.2)

    # retrieve RGB data
    params = {
        'layers': ["by_dop20c"],
        'bbox': bbox,
        'srs': srs,
        'format': format,
        'size': (width, height)
        }
    data = wms.getmap(**params, timeout=60)  # Timeout auf 60 Sekunden erhöhen
    rgb_data = np.array(Image.open(data).convert('RGB'))

    # retrieve IRG data
    params = {
        'layers': ["by_dop20cir"],
        'bbox': bbox,
        'srs': srs,
        'format': format,
        'size': (width, height)
        }
    data = wms.getmap(**params, timeout=60)  # Timeout auf 60 Sekunden erhöhen
    irg_data = np.array(Image.open(data).convert('RGB'))

    # combine datasets
    rgbi_data = np.array([rgb_data[:,:,0], rgb_data[:,:,1], rgb_data[:,:,2], irg_data[:,:,0]], dtype=np.uint8)
    # shape: [channels, height, width]

    # check image properties
    assert np.all(np.equal(rgbi_data.shape, (4, height, width)))

    # write data to geotiff file
    transform = from_bounds(*bbox, rgbi_data.shape[2], rgbi_data.shape[1])

    filename = "{:d}_{:d}_rgbi_20cm.tif".format(int(bbox[0]), int(bbox[1]))
    with rasterio.open(
                os.path.join(targetdir, filename),
                'w',
                driver='GTiff',
                height=height,
                width=width,
                count=rgbi_data.shape[0],
                dtype='uint8',
                crs=CRS.from_string(srs),
                transform=transform) as dataset:
        dataset.write(rgbi_data)

    return rgbi_data




# Download fuer Stadt Memmingen

import os

wue_bbox = (4365264, 5315180, 4365944, 5320884)
# W, S, N, E

width = 200  # px
height = 200 # px
gsd = 0.2 # m/px

os.mkdir('memmingen/') if not os.path.exists('memmingen') else None

for x in tqdm(np.arange(wue_bbox[0], wue_bbox[2], width*gsd)):
    for y in np.arange(wue_bbox[1], wue_bbox[3], height*gsd):
        
        coo = (x, y) # SW corner of tile
        bb = (x, y, x+width*gsd, y+height*gsd) # (x_min, y_min, x_max, y_max)
        
        get_wms_tile(bb, targetdir='memmingen/')





# Umwandlung der (RGBI) TIFF-Files in (RGB) PNGs
from PIL import Image

datadir = 'wuerzburg/'
targetdir = 'wuerzburg_rgb/'
filenames = os.listdir(datadir)

os.mkdir(targetdir) if not os.path.exists(targetdir) else None

for filename in filenames:
    dataset = rasterio.open(os.path.join(datadir, filename))

    img = np.dstack(dataset.read()[:3])
    image = Image.fromarray(img.astype('uint8'), 'RGB')
    image.save(os.path.join(targetdir, filename.replace('.tif', '_rgb.png')))

    print(filename)