"""export.py

Export layouts from QGIS project to pdf or jpg.

Each exported map will contain a different combination of EP/type, e.g.
0.05/'ZSA'

Usage:
    > export.bat

This script must be run with the version of python installed with QGIS, that
is why 'export.bat' is required.

D Howe
d.howe@wrl.unsw.edu.au
2022-05-17
"""
import re
import os
from qgis.core import (
    QgsApplication,
    QgsProject,
    QgsLayoutExporter,
    QgsLayoutItemLabel,
)

QGZ_PATH = 'roches.qgz'
OUTPUT_DIR = 'maps'
FORMAT = 'jpg'

# Initialise QGIS
qgs = QgsApplication([], False)

# Load project
project = QgsProject.instance()
project.read(QGZ_PATH)


def unique_attributes(layer):
    """Get unique values for year, EP, and profile"""
    layer.setSubsetString('')  # Remove filter

    # Get all available values for encounter probability and profile type
    eps = []
    types = []
    years = []
    for f in layer.getFeatures():
        eps.append(f.attribute('ep'))
        types.append(f.attribute('type'))
        years.append(f.attribute('year'))

    # Keep unique values
    eps = list(set(eps))
    types = list(set(types))
    years = list(set(years))

    return eps, types, years


# Get layers with hazard lines (assuming there is a year in the layer name)
map_layers = project.mapLayers().items()
layers = [v for k, v in map_layers if re.search(r'\d{4}', v.name())]

try:
    next(layers[0].getFeatures())
except StopIteration:
    msg = 'No hazard lines found. Check shapefile data.'
    raise ValueError(msg)


# Get all EP values and profile types from first hazard line layer
eps, types, years = unique_attributes(layers[0])

# BUG: repeat first file because of QGIS's incorrect legend placement
years = [years[0]] + years

# Get layouts
layouts = project.layoutManager().layouts()

# Create outout directory
os.makedirs(OUTPUT_DIR, exist_ok=True)

for layout in layouts:
    name = layout.name()

    # Set up plot settings
    if FORMAT == 'pdf':
        settings = QgsLayoutExporter.PdfExportSettings()
        # settings.dpi = -1  # Native resolution
        exporter = QgsLayoutExporter(layout).exportToPdf
    elif FORMAT == 'jpg':
        settings = QgsLayoutExporter.ImageExportSettings()
        exporter = QgsLayoutExporter(layout).exportToImage
    else:
        raise ValueError("'FORMAT' must be 'jpg' or 'pdf'")

    for t in types:
        for e in eps:
            for layer in layers:
                year = re.search(r'\d{4}', layer.name()).group()
                substring = f"ep='{e}' and type='{t}' and year={year}"
                layer.setSubsetString(substring)

            # Update label
            label = f'Type: {t}\nEP: {e}'
            for item in layout.items():
                if isinstance(item, QgsLayoutItemLabel):
                    item.setText(label)

            map_path = os.path.join(OUTPUT_DIR, f'{name}-{t}-EP={e}.{FORMAT}')
            print(f'Exporting {os.path.basename(map_path)}...')

            res = exporter(map_path, settings)
            if res != 0:
                msg = 'Export failed. Does output folder exist?'
                raise ValueError(msg)