# Copyright 2015 Don Drake don@drakeconsulting.com
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import sys
import os
import unittest
import datetime
import time
from copy import deepcopy
src_dir = os.path.dirname(os.path.realpath(__file__))
sys.path.append(os.path.join(src_dir, '..'))

from pyspark.sql.types import Row, StructType, StructField, IntegerType, StringType, BinaryType, BooleanType, DateType, TimestampType, DoubleType, FloatType, ByteType, LongType, ShortType
from SparkTestCase import SparkTestCase
from smartframes import SmartFrames


class SimpleTable(SmartFrames.SmartFrames):
    schema = StructType( sorted(
        [
        StructField("pk_id", IntegerType()),
        StructField("first_name", StringType()),
        ],
        key = lambda x: x.name))
    skipSelectedFields = []

class ComplexTable(SmartFrames.SmartFrames):
    schema = StructType( sorted(
        [
        StructField("pk_id", IntegerType()),
        StructField("string", StringType()),
        StructField("binary", BinaryType()),
        StructField("boolean", BooleanType()),
        StructField("date", DateType()),
        StructField("time", TimestampType()),
        StructField("double1", DoubleType()),
        StructField("double2", DoubleType()),
        StructField("float1", FloatType()),
        StructField("float2", FloatType()),
        StructField("byte", ByteType()),
        StructField("integer", IntegerType()),
        StructField("along", LongType()),
        StructField("short", ShortType()),
        ],
        key = lambda x: x.name))
    skipSelectedFields = []


class TestSmartFrames(SparkTestCase):

    def testSimpleTable(self):
        simpleTable = SimpleTable()
        self.assertEquals(simpleTable.schema, SimpleTable().schema)

        s1 = SimpleTable()
        s1.pk_id = 1
        s1.first_name = 'Don'
                        
        s2 = SimpleTable()
        s2.pk_id = 2
        s2.first_name = 'Dan'
        df = self.sqlCtx.createDataFrame(self.sc.parallelize([s1.createRow(), s2.createRow()]), s1.schema)
        self.assertEquals(2, df.count())

        fileName = self.tempdir + '/simple.table'
        df.saveAsParquetFile(fileName)

        df2 = self.sqlCtx.parquetFile(fileName)
        self.assertEquals(sorted(df.collect()), sorted(df2.collect()))


    def testComplexTable(self):
        complexTable = ComplexTable()
        self.assertEquals(complexTable.schema, ComplexTable().schema)

        s1 = ComplexTable()
        s1.pk_id = 1
        s1.string = 'abcdefghijklmnopqrstuvwxyz'
        s1.binary = bytearray(b"0xDEADBEEF")
        s1.boolean = True
        s1.date = datetime.date(2015, 10, 3)
        s1.time = datetime.datetime(2015, 10, 3, 14, 33)
        s1.double1 = 1
        s1.double2 = 2.2
        s1.float1 = 1
        s1.float2 = 2.2
        s1.byte = 100
        s1.integer = 10000
        s1.along = 10000
        s1.short = 10

        df = self.sqlCtx.createDataFrame(self.sc.parallelize([s1.createRow()]), s1.schema)
        fileName = self.tempdir + '/complex.table'
        df.saveAsParquetFile(fileName)

        df2 = self.sqlCtx.parquetFile(fileName)
        self.assertEquals(sorted(df.collect()), sorted(df2.collect()))
        r1 = df2.collect()[0]
        print "r1=", r1
        self.assertEquals(r1.pk_id, s1.pk_id)
        self.assertEquals(r1.string, s1.string)
        self.assertEquals(r1.binary, s1.binary)
        self.assertEquals(r1.boolean, s1.boolean)
        self.assertEquals(r1.date, s1.date)
        self.assertEquals(r1.time, s1.time)
        self.assertEquals(r1.double1, s1.double1)
        self.assertEquals(r1.double2, s1.double2)
        self.assertEquals(r1.float1, s1.float1)
        # AssertionError: 2.200000047683716 != 2.2
        #self.assertEquals(r1.float2, s1.float2)
        self.assertEquals(r1.byte, s1.byte)
        self.assertEquals(r1.integer, s1.integer)
        self.assertEquals(r1.along, s1.along)
        self.assertEquals(r1.short, s1.short)

    def testComplexTableTiming(self):
        s1 = ComplexTable()
        s1.pk_id = 1
        s1.string = 'abcdefghijklmnopqrstuvwxyz'
        s1.binary = bytearray(b"0xDEADBEEF")
        s1.boolean = True
        s1.date = datetime.date(2015, 10, 3)
        s1.time = datetime.datetime(2015, 10, 3, 14, 33)
        s1.double1 = 1
        s1.double2 = 2.2
        s1.float1 = 1
        s1.float2 = 2.2
        s1.byte = 100
        s1.integer = 10000
        s1.along = 10000
        s1.short = 10

        numRows = 10000
        rows = []
        start = time.clock()
        for n in xrange(numRows):
            

        end = time.clock()
        print "Duration for ", numRows, " is ", (end - start)
        df = self.sqlCtx.createDataFrame(self.sc.parallelize(rows), s1.schema)
        # duration for v.1.0.1 ~ 2.4 seconds for 10000 rows
        # duration for v.1.1.0 ~ 1.6 seconds for 10000 rows
        count = df.count()
        self.assertEquals(count, numRows)

if __name__ == '__main__':
    unittest.main()

rows.append(deepcopy(s1.createRow()))