﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using IronPython.Runtime;
using System.Collections;

namespace IronMath
{
    public class ndarray : IEnumerable
    {

        private int size;
        private int row;
        private int col;
        private List<int> shape;
        private int length;
        private List<object> tmpitems;

        public ndarray(List data)
        {
            //from the data figure out the dimension
            shape = new List<int>();
            GetShapeFromList(data);
            SetLength();
            tmpitems = new List<object>();
            UnpackData(data);
            if (tmpitems.Count != length)
                throw new ArgumentException();
        }

        public ndarray(IEnumerable<object> data, PythonTuple shape)
        {
            tmpitems = new List<object>(data);
            this.shape = new List<int>();
            length = 1;
            foreach (int axis in shape)
            {
                length *= axis;
                this.shape.Add(axis);
            }
            if (tmpitems.Count != length)
                throw new ArgumentException();
        }

        internal ndarray(List<int> shape, List<object> data)
        {
            this.shape = shape;
            this.tmpitems = data;
            this.length = data.Count;
        }

        private void UnpackData(List data)
        {
            for (int i = 0; i < data.Count; i++)
                if (data[i].GetType() == typeof(List))
                    UnpackData((List)data[i]);
                else
                    tmpitems.Add(data[i]);
        }

        private void GetShapeFromList(List data)
        {
            shape.Add(data.Count);
            if (data[0].GetType() == typeof(List))
            {
                GetShapeFromList((List)data[0]);
            }
        }

        private void SetLength()
        {
            length = 1;
            foreach (var axis in shape)
            {
                length *= axis;
            }
        }

        public int Length
        {
            get { return length; }
        }

        public PythonTuple Shape
        {
            get { return new PythonTuple(this.shape.ToArray()); }
            set { SetShape(value); }
        }

        private void SetShape(PythonTuple shape)
        {
            int len = 1;
            List<int> newShape = new List<int>();
            foreach (int i in shape)
            {
                len *= i;
                newShape.Add(i);
            }

            if (len != length)
                throw new ArgumentException("Not a valid shape");
            this.shape = newShape;
        }

        IEnumerator IEnumerable.GetEnumerator()
        {
            return tmpitems.GetEnumerator();
        }

        public object __getitem__(params object[] index)
        {
            List<object> data = new List<object>(index);
            if (data.Count < shape.Count)
            {
                //sub ndarray
                //pad it with empty ':'
                int start = data.Count;
                for (int i = start; i < shape.Count; i++)
                    data.Add(new Slice(null));
            }
            List<IEnumerable<int>> indices = new List<IEnumerable<int>>();
            for (int i = 0; i < data.Count; i++)
                indices.Add(Indices(data[i], i + 1).Distinct());

            //get the new shape
            List<int> newShape = new List<int>();
            for (int i = 0; i < indices.Count; i++)
                newShape.Add(indices[i].Count());

            //check to see if there are multi-dimensional indexes
            while (indices.Count > 1)
            {
                var setA = indices[0];
                var setB = indices[1];
                indices.RemoveRange(0, 2);
                indices.Insert(0, CombineIndex(setA, setB));
            }

            if (indices[0].Count() == 1)
                return tmpitems[indices[0].ToList()[0]];

            var sliceData = from i in indices[0].ToList()
                            select tmpitems[i];

            ndarray ret = new ndarray(newShape, sliceData.ToList());
            return ret;
        }



        private int DimensionOffset(int dimension)
        {
            int off = 1;
            for (int i = dimension; i < shape.Count; i++)
                off *= shape[i];
            return off;
        }

        private IEnumerable<int> Indices(object selection, int dimension)
        {
            List<int> ind = new List<int>();
            int baseOffset = DimensionOffset(dimension);
            //check the type of the selection
            if (selection.GetType() == typeof(List))
            {
                //its a list and can contain aditional selection objects
                foreach (object sel in (List)selection)
                {
                    ind.AddRange(Indices(sel, dimension));
                }
            }
            else if (selection.GetType() == typeof(Slice))
            {
                //its a python slice operator
                ndslice sel = new ndslice((Slice)selection, shape[dimension - 1]);
                foreach (int i in sel.GetIndices())
                    ind.Add(i * baseOffset);
            }
            else
            {
                //could be just a number
                int i = (int)selection;
                ind.Add(i * baseOffset);
            }

            return ind;
        }

        private IEnumerable<int> CombineIndex(IEnumerable<int> a, IEnumerable<int> b)
        {
            //var result = a.SelectMany(x => b.Select(y => x + y));
            var result = from x in a
                         from y in b
                         select (x + y);

            return result;
        }


        public ndarray __add__(int value)
        {
            List<object> tmp = new List<object>();
            for (int i = 0; i < length; i++)
                tmp.Add((int)tmpitems[i] + value);
            return new ndarray(shape, tmp);
        }


        public ndarray __add__(double value)
        {
            List<object> tmp = new List<object>();
            for (int i = 0; i < length; i++)
                tmp.Add(Convert.ToDouble(tmpitems[i]) + value);
            return new ndarray(shape, tmp);
        }

        public ndarray __add__(long value)
        {
            List<object> tmp = new List<object>();
            for (int i = 0; i < length; i++)
                tmp.Add((long)tmpitems[i] + value);
            return new ndarray(shape, tmp);
        }


        public int __len__()
        {
            return length;
        }


        public string __str__()
        {
            return GenerateString(0, 1);
        }

        private string GenerateString(int offset, int dimension)
        {
            int baseOffset = DimensionOffset(dimension);
            StringBuilder sb = new StringBuilder();
            List<string> tmp = new List<string>();
            //check if this is the last dimension
            if (dimension < shape.Count)
            {
                sb.Append("[");
                //call the other level
                for (int i = 0; i < shape[dimension - 1]; i++)
                    tmp.Add(GenerateString(i * baseOffset + offset, dimension + 1));

                sb.Append(string.Join(",\n", tmp.ToArray()));
                sb.Append("]");
            }
            else
            {
                sb.Append("[");
                for (int i = 0; i < shape[dimension - 1]; i++)
                    tmp.Add(tmpitems[i * baseOffset + offset].ToString());

                sb.Append(string.Join(",\t", tmp.ToArray()));
                sb.Append("]");
            }
            return sb.ToString();
        }



        public List flatten()
        {
            List l = new List();
            l.extend(tmpitems);
            return l;
        }

        internal class ndslice
        {
            int start;
            int stop;
            int step;
            public ndslice(Slice slice, int length)
            {
                start = (slice.start == null) ? 0 : (int)slice.start;
                stop = (slice.stop == null) ? length : (int)slice.stop;
                step = (slice.step == null) ? 1 : (int)slice.step;

                if (start < 0) start += length;
                if (stop < 0) stop += length;
                if (stop > length) stop = length; //or throw an exception??
            }

            public int Start { get { return start; } }
            public int Stop { get { return stop; } }
            public int Step { get { return step; } }

            public IEnumerable GetIndices()
            {
                for (int i = start; i < stop; i+= step)
                    yield return i;
            }
        }


    }
}
