using System.Collections.Generic;
using System.IO;
using System;
using System.Linq;
using UnityEditor;
using UnityEngine;
using System.Text.RegularExpressions;

static class Comparer
{
    public static IEnumerable<T> OrderByNatural<T>(this IEnumerable<T> items, Func<T, string> selector, StringComparer stringComparer = null)
    {
        var regex = new Regex(@"\d+", RegexOptions.Compiled);

        int maxDigits = items
                      .SelectMany(i => regex.Matches(selector(i)).Cast<Match>().Select(digitChunk => (int?)digitChunk.Value.Length))
                      .Max() ?? 0;

        return items.OrderBy(i => regex.Replace(selector(i), match => match.Value.PadLeft(maxDigits, '0')), stringComparer ?? StringComparer.CurrentCulture);
    }
}

[CustomEditor(typeof(CubeManager))]
public class CubeManagerEditor : Editor
{

    public string importPath = "C:\\Users\\admin_vyskoc65\\Desktop\\data\\N";
    public int limit = 0;
    public int stepSize = 1;
    public bool recursiveImport = false;


    private CubeManager _cubeManager;
    private CubeFileInfo _cubeFileInfo;

    static void Init()
    {
        var window = EditorWindow.GetWindow(typeof(CubeManagerEditor));
        window.Show();
    }

    public override void OnInspectorGUI()
    {
        DrawDefaultInspector();

        EditorGUILayout.LabelField("", GUI.skin.horizontalSlider);

        importPath = EditorGUILayout.TextField("Import Path", importPath);
        limit = EditorGUILayout.IntField("Limit Number of Files", limit);
        stepSize = EditorGUILayout.IntField("Step of Imported Frames", stepSize);

        _cubeManager = (CubeManager)target;

        if (GUILayout.Button("Import Files"))
        {
            ImportFiles();
        }
    }


    public void ImportFiles()
    {
        if (!AssetDatabase.IsValidFolder(Path.Combine("Assets/Resources", _cubeManager.datasetPrefix)))
        {
            AssetDatabase.CreateFolder("Assets/Resources", _cubeManager.datasetPrefix);
        }

        if (Directory.Exists(importPath))
        {
            ProcessDirectory(importPath);
        }
        else
        {
            Debug.LogError(importPath + " is not a valid file or directory.");
        }
    }


    public void ProcessDirectory(string importDirectory)
    {
        string texturePrefixPath = _cubeManager.asAsset ? Path.Combine("Assets/Resources", _cubeManager.datasetPrefix) : Path.Combine(Application.dataPath, "..", _cubeManager.datasetPrefix);
        string outputPath = Path.Combine(texturePrefixPath, _cubeManager.datasetName);
        if (!_cubeManager.asAsset) { 
            outputPath = Path.GetFullPath(outputPath);  // normalize to canonical path
        }
        string extension = _cubeManager.asAsset ? ".asset" : ".texture";

        if (_cubeManager.asAsset)
        {
            if (!AssetDatabase.IsValidFolder(outputPath))
            {
                AssetDatabase.CreateFolder(texturePrefixPath, _cubeManager.datasetName);
            }
        }
        else
        {
            Directory.CreateDirectory(outputPath);
        }

        _cubeFileInfo = new CubeFileInfo();

        int readCounter = 0;
        int skipCounter = 0;
        List<string> fileEntries = Directory.GetFiles(importDirectory).ToList().OrderByNatural(f => f).ToList();
        //var fileEntries = OrderByAlphaNumeric(Directory.GetFiles(importDirectory));
        float numFiles = fileEntries.Count;
        foreach (string fileName in fileEntries)
        {
            if (skipCounter % stepSize == 0){
                if (limit == 0 || limit >= readCounter)
                {
                    EditorUtility.DisplayProgressBar("Importing Data Files", "Importing " + fileName, (float)skipCounter / numFiles);
                    string suffix = readCounter.ToString("D6") + extension;
                    string volumeFileName = "volume_" + suffix;
                    string positionsFileName = "positions_" + suffix;
                    string volumeFile = outputPath + "/" + volumeFileName;
                    string positionsFile = outputPath + "/" + positionsFileName;
                    CubeDomainDetails details = ProcessFile(fileName, volumeFile, positionsFile);
                    if (readCounter == 0) // Set domain details from the first file only
                    {
                        _cubeFileInfo.origin = details.origin;
                        _cubeFileInfo.size = details.size;
                        _cubeFileInfo.nAtoms = details.nAtoms;
                        _cubeFileInfo.axes.Add(details.axes[0]);
                        _cubeFileInfo.axes.Add(details.axes[1]);
                        _cubeFileInfo.axes.Add(details.axes[2]);
                    }


                    _cubeFileInfo.ids.Add(readCounter);
                    _cubeFileInfo.volumes.Add(volumeFileName);
                    _cubeFileInfo.positions.Add(positionsFileName);
                    _cubeFileInfo.norms.Add(details.normalization);
                    ++readCounter;
                    if (readCounter % 100 == 0){
                        AssetDatabase.SaveAssets();
                        EditorUtility.UnloadUnusedAssetsImmediate();
                        GC.Collect();
                    }
                }
            }
            ++skipCounter;
        }
        EditorUtility.ClearProgressBar();

        if (recursiveImport)
        {
            Debug.LogError("recursive import not implemented");
            return;
            /*
            string[] subdirectoryEntries = Directory.GetDirectories(targetDirectory).OrderBy(f => f);
            foreach (string subdirectory in subdirectoryEntries)
            {
                ProcessDirectory(subdirectory);
            }
            */
        }

        string json = JsonUtility.ToJson(_cubeFileInfo);
        string jsonFileName = outputPath + "/fileInfo.json";
        File.WriteAllText(jsonFileName, json);
    }


    // Reads Gaussian Cube files http://paulbourke.net/dataformats/cube/
    public CubeDomainDetails ProcessFile(string inputFileName, string volumeFileName, string positionsFileName)
    {
        //Debug.Log("Processing file '" + inputFileName + "'...");

        StreamReader inputFile = new StreamReader(inputFileName);

        CubeDomainDetails details = new CubeDomainDetails();

        //  The first two lines of the header are comments, they are generally ignored by parsing packages or used as two default labels. 
        inputFile.ReadLine();
        inputFile.ReadLine();

        // The third line has the number of atoms included in the file followed by the position of the origin of the volumetric data. 
        string line = inputFile.ReadLine();
        var items = line.Split(new char[0], StringSplitOptions.RemoveEmptyEntries);
        
        int nAtoms = int.Parse(items[0]);
        details.nAtoms = nAtoms;
        
        Vector3 volumeOrigin = new Vector3(NumParse.Float(items[1]), NumParse.Float(items[2]), NumParse.Float(items[3]));
        details.origin = volumeOrigin;

        /* The next three lines give the number of voxels along each axis (x, y, z) followed by the axis vector. 
           Note this means the volume need not be aligned with the coordinate axis, indeed it also means it may 
           be sheared although most volumetric packages won't support that. The length of each vector is the length 
           of the side of the voxel thus allowing non cubic volumes. If the sign of the number of voxels in 
           a dimension is positive then the units are Bohr, if negative then Angstroms.
        */
        for (int i = 0; i < 3; ++i)
        {
            line = inputFile.ReadLine();
            items = line.Split(new char[0], StringSplitOptions.RemoveEmptyEntries);
            details.size[i] = int.Parse(items[0]);
            details.axes[i] = new Vector3(NumParse.Float(items[1]), NumParse.Float(items[2]), NumParse.Float(items[3]));
        }

        Vector3 physicalExtent = new Vector3(
            Math.Abs(details.size[0] * details.axes[0].magnitude),
            Math.Abs(details.size[1] * details.axes[1].magnitude),
            Math.Abs(details.size[2] * details.axes[2].magnitude)
            );


        /* The last section in the header is one line for each atom consisting of 5 numbers, the first is the atom 
           number, the second is the charge, and the last three are the x,y,z coordinates of the atom center. 
        */
        var positionTexture = new Texture2D(nAtoms, 1, TextureFormat.RGBAFloat, false);
        var positions = new Color[nAtoms];
        for (int i = 0; i < nAtoms; ++i)
        {
            line = inputFile.ReadLine();
            items = line.Split(new char[0], StringSplitOptions.RemoveEmptyEntries);
            positions[i] = new Color(NumParse.Float(items[2])/physicalExtent[0], NumParse.Float(items[3])/physicalExtent[1], NumParse.Float(items[4])/physicalExtent[2], NumParse.Float(items[1])/300f);
        }
        positionTexture.filterMode = FilterMode.Point;
        positionTexture.SetPixels(positions);
        positionTexture.Apply();

        if (_cubeManager.asAsset)
        {
            AssetDatabase.CreateAsset(positionTexture, positionsFileName);
        }
        else
        {
            byte[] rawData = positionTexture.GetRawTextureData<byte>().ToArray();
            File.WriteAllBytes(positionsFileName, rawData);
        }

        /* The volumetric data is stored as a sequence of floats iterating as x,y,z. */
        line = inputFile.ReadLine();
        items = line.Split(new char[0], StringSplitOptions.RemoveEmptyEntries);
        int lineElementCount = 0;

        float[,,] values = new float[details.size[0], details.size[1], details.size[2]];

        float maxValue = float.NegativeInfinity;
        float minValue = float.PositiveInfinity;

        for (int x = 0; x < details.size[0]; x++)
        {
            for (int y = 0; y < details.size[1]; y++)
            {
                for (int z = 0; z < details.size[2]; z++)
                {
                    // Go to next line
                    if (lineElementCount >= items.Length)
                    {
                        line = inputFile.ReadLine();
                        items = line.Split(new char[0], StringSplitOptions.RemoveEmptyEntries);
                        lineElementCount = 0;
                    }

                    float value = NumParse.Float(items[lineElementCount]);
                    if (value > maxValue) maxValue = value;
                    if (value < minValue) minValue = value;

                    values[x, y, z] = value;

                    ++lineElementCount;
                }
            }
        }

        details.normalization = Math.Max(Math.Abs(minValue), Math.Abs(maxValue));
        float normFactor = 1.0f / details.normalization;

        TextureFormat format = TextureFormat.RFloat;
        TextureWrapMode wrapMode = TextureWrapMode.Clamp;

        Texture3D texture = new Texture3D(details.size[0], details.size[1], details.size[2], format, false);
        texture.wrapMode = wrapMode;

        Color[] colors = new Color[details.size[0] * details.size[1] * details.size[2]];

        for (int x = 0; x < details.size[0]; x++)
        {
            for (int y = 0; y < details.size[1]; y++)
            {
                int yOffset = y * details.size[0];
                for (int z = 0; z < details.size[2]; z++)
                {
                    int zOffset = z * details.size[0] * details.size[1];
                    colors[x + yOffset + zOffset] = new Color(normFactor * values[x, y, z], 0.0f, 0.0f, 1.0f);
                }
            }
        }

        texture.SetPixels(colors);
        texture.Apply();
        
        if (_cubeManager.asAsset)
        {
            AssetDatabase.CreateAsset(texture, volumeFileName);
        }
        else
        {
            byte[] rawData = texture.GetPixelData<byte>(0).ToArray();
            File.WriteAllBytes(volumeFileName, rawData);
        }

        return details;
    }

}
