import argparse
import os
import xml.etree.ElementTree as et
from zipfile import ZipFile

import numpy as np


def main():
    print("FILENAME :", zipFileName)
    baseName, _ = os.path.splitext(zipFileName)
    print("DIRNAME  :", baseName)
    os.makedirs(baseName, exist_ok=True)
    existing_zip= ZipFile(zipFileName)
    existing_zip.extractall(baseName)
    existing_zip.close()

    xNum = 225
    zNum = 150

    data = np.full((zNum * 10, xNum * 10), np.nan, dtype=np.float)

    for fileName in os.listdir(baseName):
        _, ext = os.path.splitext(fileName)
        if ext == ".xml":
            print("FILENAME :", fileName)
            tree = et.parse(baseName + "//" + fileName)
            root = tree.getroot()
            mesh = root[2][6].text
            N = (9 - int(mesh[6])) * zNum
            E = int(mesh[7]) * xNum
            print(
                "LOCATION :", mesh, root[2][7][0][0][0].text, root[2][7][0][0][1].text
            )

            curData = np.full((zNum, xNum), np.nan, dtype=np.float)
            rawData = root[2][7][2][0][1].text.split()
            xStart, zStart = root[2][7][3][0][1].text.split()
            xStart = int(xStart)
            zStart = int(zStart)
            x=xStart
            z=zStart
            for raw in rawData:
                if raw.split(",")[0] == "地表面":
                    curData[z][x]=float(raw.split(",")[1])
                x+=1
                if x >= xNum:
                    x=0
                    z+=1
                    if z >= zNum:
                        break
            data[N : N + zNum, E : E + xNum] = curData
            print("MAINDATA :", curData)
        os.remove(baseName + "//" + fileName)

    os.removedirs(baseName)

    data *= yScale*20/(np.nanmax(data)-np.nanmin(data))
    lowestLayer = np.nanmin(data)
    data[np.isnan(data)] = lowestLayer - 3
    data -= (lowestLayer - 5)

    file = open(baseName + ".obj", "w")

    for z in range(zNum * xzScale1):
        print("\rPROCESSING :{0}".format(int(100*z/(zNum * xzScale1))), end="")
        v = ""
        f = ""
        for x in range(xNum * xzScale1):
            curData = (
                str(
                    data[
                        z * xzScale2 : z * xzScale2 + xzScale2 - 1,
                        x * xzScale2 : x * xzScale2 + xzScale2 - 1,
                    ].mean()
                )
            )
            curPos = (z * xNum * xzScale1 + x) * 8
            v += (
                "v " + str(x) + " " + curData + " " + str(z) + "\n"
                "v " + str(x + 1) + " " + curData + " " + str(z) + "\n"
                "v " + str(x) + " " + curData + " " + str(z + 1) + "\n"
                "v " + str(x + 1) + " " + curData + " " + str(z + 1) + "\n"
                "v " + str(x) + " 0 " + str(z) + "\n"
                "v " + str(x + 1) + " 0 " + str(z) + "\n"
                "v " + str(x) + " 0 " + str(z + 1) + "\n"
                "v " + str(x + 1) + " 0 " + str(z + 1) + "\n"
            )
            f += (
                "f "
                + str(curPos + 1)
                + " "
                + str(curPos + 2)
                + " "
                + str(curPos + 3)
                + "\n"
                + "f "
                + str(curPos + 2)
                + " "
                + str(curPos + 3)
                + " "
                + str(curPos + 4)
                + "\n"

                + "f "
                + str(curPos + 1)
                + " "
                + str(curPos + 2)
                + " "
                + str(curPos + 5)
                + "\n"
                + "f "
                + str(curPos + 2)
                + " "
                + str(curPos + 5)
                + " "
                + str(curPos + 6)
                + "\n"

                + "f "
                + str(curPos + 2)
                + " "
                + str(curPos + 3)
                + " "
                + str(curPos + 6)
                + "\n"
                + "f "
                + str(curPos + 3)
                + " "
                + str(curPos + 6)
                + " "
                + str(curPos + 7)
                + "\n"

                + "f "
                + str(curPos + 3)
                + " "
                + str(curPos + 4)
                + " "
                + str(curPos + 7)
                + "\n"
                + "f "
                + str(curPos + 4)
                + " "
                + str(curPos + 7)
                + " "
                + str(curPos + 8)
                + "\n"

                + "f "
                + str(curPos + 4)
                + " "
                + str(curPos + 1)
                + " "
                + str(curPos + 8)
                + "\n"
                + "f "
                + str(curPos + 1)
                + " "
                + str(curPos + 8)
                + " "
                + str(curPos + 2)
                + "\n"
            )

        file.write(v + f)

    file.write(
        "v 0 0 0\n"
        + "v " + str(xNum * xzScale1) + " 0 0\n"
        + "v 0 0 " + str(zNum * xzScale1) + "\n"
        + "v " + str(xNum * xzScale1) + " 0 " + str(zNum * xzScale1) + "\n"
        + "f"
        + str(zNum * xzScale1 * xNum * xzScale1 * 8 + 1)
        + " "
        + str(zNum * xzScale1 * xNum * xzScale1 * 8 + 2)
        + " "
        + str(zNum * xzScale1 * xNum * xzScale1 * 8 + 3)
        + "\n"
        + "f"
        + str(zNum * xzScale1 * xNum * xzScale1 * 8 + 2)
        + " "
        + str(zNum * xzScale1 * xNum * xzScale1 * 8 + 3)
        + " "
        + str(zNum * xzScale1 * xNum * xzScale1 * 8 + 4)
        + "\n"
    )

    file.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("filename", help="Specify a DEM ZIP file")
    parser.add_argument('-x', '--xzscale', type=int, default=2, help="Specify the horizontal scale with a value between 1 and 4. The higher the number, the more precise.")
    parser.add_argument('-y', '--yscale', type=float, default=1.0, help="Specify the vertical scale.")
    args = parser.parse_args()

    yScale = args.yscale
    zipFileName=args.filename

    # scale1 * scale2 == 10
    xzScale = args.xzscale
    if xzScale == 1:
        xzScale1 = 1
        xzScale2 = 10
    elif xzScale == 2:
        xzScale1 = 2
        xzScale2 = 5
    elif xzScale == 3:
        xzScale1 = 5
        xzScale2 = 2
    else:
        xzScale1 = 10
        xzScale2 = 1

    main()